Skip to content

Commit

Permalink
fix function signature
Browse files Browse the repository at this point in the history
  • Loading branch information
alisterburt committed Feb 22, 2024
1 parent 4128727 commit 3d3fc09
Showing 1 changed file with 41 additions and 9 deletions.
50 changes: 41 additions & 9 deletions src/starfile/functions.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from __future__ import annotations

from typing import TYPE_CHECKING, Dict, List, Union
from typing import TYPE_CHECKING, Dict, List, Union, Optional

if TYPE_CHECKING:
import pandas as pd
Expand All @@ -15,14 +15,31 @@
from os import PathLike


def read(filename: PathLike, read_n_blocks: int = None, always_dict: bool = False, parse_as_string: List[str] = None):
"""
Read a star file into a pandas dataframe or dict of pandas dataframes
def read(
filename: PathLike,
read_n_blocks: Optional[int] = None,
always_dict: bool = False,
parse_as_string: List[str] = []
) -> Union[DataBlock, Dict[DataBlock]]:
"""Read data from a STAR file.
default behaviour in the case of only one data block being present in the STAR file is to
return only a dataframe, this can be changed by setting 'always_dict=True'
"""
Basic data blocks are read as dictionaries. Loop blocks are read as pandas
dataframes. When multiple data blocks are present a dictionary of datablocks is
returned. When a single datablock is present only the block is returned by default.
To force returning a dectionary even when only one datablock is present set
`always_dict=True`.
Parameters
----------
filename: PathLike
File from which to read data.
read_n_blocks: int | None
Limit reading the file to the first n data blocks.
always_dict: bool
Always return a dictionary, even when only a single data block is present.
parse_as_string: list[str]
A list of keys or column names which will not be coerced to numeric values.
"""
parser = StarParser(filename, n_blocks_to_read=read_n_blocks, parse_as_string=parse_as_string)
if len(parser.data_blocks) == 1 and always_dict is False:
return list(parser.data_blocks.values())[0]
Expand All @@ -38,9 +55,24 @@ def write(
na_rep: str = '<NA>',
quote_character: str = '"',
quote_all_strings: bool = False,
**kwargs,
**kwargs
):
"""Write data blocks as STAR files."""
"""Write data to disk in the STAR format.
Parameters
----------
data: DataBlock | Dict[str, DataBlock] | List[DataBlock]
Data to be saved to file. DataBlocks are dictionaries or dataframes.
If a dictionary of datablocks are passed the keys will be the data block names.
filename: PathLike
Path where the file will be saved.
float_format: str
Float format string which will be passed to pandas.
sep: str
Separator between values, will be passed to pandas.
na_rep: str
Representation of null values, will be passed to pandas.
"""
StarWriter(
data,
filename=filename,
Expand Down

0 comments on commit 3d3fc09

Please sign in to comment.