diff --git a/csvy/__init__.py b/csvy/__init__.py index c20172d..9955d89 100644 --- a/csvy/__init__.py +++ b/csvy/__init__.py @@ -1,6 +1,4 @@ -""" -Python reader/writer for CSV files with YAML header information. -""" +"""Python reader/writer for CSV files with YAML header information.""" __version__ = "0.2.2" from .readers import ( # noqa: F401 diff --git a/csvy/readers.py b/csvy/readers.py index e51bcd9..ba9803e 100644 --- a/csvy/readers.py +++ b/csvy/readers.py @@ -1,6 +1,8 @@ +"""A collection of functions for parsing CSVY files.""" + import logging from pathlib import Path -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any import yaml @@ -45,8 +47,8 @@ def get_comment(line: str, marker: str = "---") -> str: def read_header( - filename: Union[Path, str], marker: str = "---", **kwargs: Any -) -> Tuple[Dict[str, Any], int, str]: + filename: Path | str, marker: str = "---", **kwargs: Any +) -> tuple[dict[str, Any], int, str]: """Read the yaml-formatted header from a file. Args: @@ -80,8 +82,8 @@ def read_header( def read_metadata( - filename: Union[Path, str], marker: str = "---", **kwargs: Any -) -> Dict[str, Any]: + filename: Path | str, marker: str = "---", **kwargs: Any +) -> dict[str, Any]: """Read the yaml-formatted metadata from a file. Args: @@ -96,11 +98,11 @@ def read_metadata( def read_to_array( - filename: Union[Path, str], + filename: Path | str, marker: str = "---", - csv_options: Optional[Dict[str, Any]] = None, - yaml_options: Optional[Dict[str, Any]] = None, -) -> Tuple[NDArray, Dict[str, Any]]: + csv_options: dict[str, Any] | None = None, + yaml_options: dict[str, Any] | None = None, +) -> tuple[NDArray, dict[str, Any]]: """Reads a CSVY file into dict with the header and array with the data. Args: @@ -131,11 +133,11 @@ def read_to_array( def read_to_dataframe( - filename: Union[Path, str], + filename: Path | str, marker: str = "---", - csv_options: Optional[Dict[str, Any]] = None, - yaml_options: Optional[Dict[str, Any]] = None, -) -> Tuple[DataFrame, Dict[str, Any]]: + csv_options: dict[str, Any] | None = None, + yaml_options: dict[str, Any] | None = None, +) -> tuple[DataFrame, dict[str, Any]]: """Reads a CSVY file into dict with the header and a DataFrame with the data. Possible 'skiprows' and 'comment' argument provided in the 'csv_options' dictionary @@ -169,11 +171,11 @@ def read_to_dataframe( def read_to_list( - filename: Union[Path, str], + filename: Path | str, marker: str = "---", - csv_options: Optional[Dict[str, Any]] = None, - yaml_options: Optional[Dict[str, Any]] = None, -) -> Tuple[List[List], Dict[str, Any]]: + csv_options: dict[str, Any] | None = None, + yaml_options: dict[str, Any] | None = None, +) -> tuple[list[list], dict[str, Any]]: """Reads a CSVY file into a list with the header and a nested list with the data. Args: @@ -196,7 +198,7 @@ def read_to_list( options = csv_options.copy() if csv_options is not None else {} data = [] - with open(filename, "r", newline="") as csvfile: + with open(filename, newline="") as csvfile: csvreader = csv.reader(csvfile, **options) for _ in range(nlines): diff --git a/csvy/writers.py b/csvy/writers.py index 68234cd..60584ec 100644 --- a/csvy/writers.py +++ b/csvy/writers.py @@ -1,29 +1,40 @@ +"""A collection of functions for writing CSVY files.""" + from __future__ import annotations import csv import logging +from collections.abc import Callable, Iterable from io import TextIOBase from pathlib import Path -from typing import Any, Callable, Dict, Iterable, List, Optional, Union +from typing import Any import yaml -KNOWN_WRITERS: List[Callable[[Union[Path, str], Any, str], bool]] = [] +KNOWN_WRITERS: list[Callable[[Path | str, Any, str], bool]] = [] + +def register_writer(fun: Callable[[Path | str, Any, str], bool]) -> Callable: + """Register a file writer. + + Args: + fun (Callable): The writer function. -def register_writer(fun: Callable[[Union[Path, str], Any, str], bool]) -> Callable: + Returns: + Callable: the writer function. + """ if fun not in KNOWN_WRITERS: KNOWN_WRITERS.append(fun) return fun def write( - filename: Union[Path, str], + filename: Path | str, data: Any, - header: Dict[str, Any], + header: dict[str, Any], comment: str = "", - csv_options: Optional[Dict[str, Any]] = None, - yaml_options: Optional[Dict[str, Any]] = None, + csv_options: dict[str, Any] | None = None, + yaml_options: dict[str, Any] | None = None, ) -> None: """Writes the data and header in a CSV file, formating the header as yaml. @@ -54,11 +65,11 @@ class Writer: def __init__( self, - filename: Union[Path, str], - header: Dict[str, Any], + filename: Path | str, + header: dict[str, Any], comment: str = "", - csv_options: Optional[Dict[str, Any]] = None, - yaml_options: Optional[Dict[str, Any]] = None, + csv_options: dict[str, Any] | None = None, + yaml_options: dict[str, Any] | None = None, line_buffering: bool = False, ) -> None: """Create a new Writer. @@ -72,7 +83,6 @@ def __init__( writing the header. line_buffering: Line buffering instead of chunk buffering (default False). """ - if not csv_options: csv_options = {} if not yaml_options: @@ -88,9 +98,11 @@ def __init__( self._writer = csv.writer(self._file, **csv_options) def __enter__(self) -> Writer: + """Enter the context manager.""" return self def __exit__(self, *_: Any) -> None: + """Exit the context manager.""" self._file.close() def close(self) -> None: @@ -107,8 +119,8 @@ def writerows(self, rows: Iterable[Iterable[Any]]) -> None: def write_header( - file: Union[Path, str, TextIOBase], - header: Dict[str, Any], + file: Path | str | TextIOBase, + header: dict[str, Any], comment: str = "", **kwargs: Any, ) -> None: @@ -137,7 +149,7 @@ def write_header( def write_data( - filename: Union[Path, str], data: Any, comment: str = "", **kwargs: Any + filename: Path | str, data: Any, comment: str = "", **kwargs: Any ) -> None: """Writes the tabular data to the chosen file, adding it after the header. @@ -160,7 +172,7 @@ def write_data( @register_writer def write_numpy( - filename: Union[Path, str], data: Any, comment: str = "", **kwargs: Any + filename: Path | str, data: Any, comment: str = "", **kwargs: Any ) -> bool: """Writes the numpy array to the chosen file, adding it after the header. @@ -193,7 +205,7 @@ def write_numpy( @register_writer def write_pandas( - filename: Union[Path, str], data: Any, comment: str = "", **kwargs: Any + filename: Path | str, data: Any, comment: str = "", **kwargs: Any ) -> bool: """Writes the pandas dataframe to the chosen file, adding it after the header. @@ -224,7 +236,7 @@ def write_pandas( def write_csv( - filename: Union[Path, str], data: Any, comment: str = "", **kwargs: Any + filename: Path | str, data: Any, comment: str = "", **kwargs: Any ) -> bool: """Writes the tabular to the chosen file, adding it after the header. @@ -239,7 +251,6 @@ def write_csv( Returns: True if the writer worked, False otherwise. """ - with open(filename, "a", newline="") as f: writer = csv.writer(f, **kwargs) for row in data: