From 3d3fc0972d665de4746793c05220dd038dff8adc Mon Sep 17 00:00:00 2001
From: Alister Burt <alisterburt@gmail.com>
Date: Wed, 21 Feb 2024 20:53:09 -0800
Subject: [PATCH] fix function signature

---
 src/starfile/functions.py | 50 ++++++++++++++++++++++++++++++++-------
 1 file changed, 41 insertions(+), 9 deletions(-)

diff --git a/src/starfile/functions.py b/src/starfile/functions.py
index 537d498..b563f8a 100644
--- a/src/starfile/functions.py
+++ b/src/starfile/functions.py
@@ -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
@@ -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]
@@ -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,