cleaned up

Author: czaloom

src/valor_lite/cache.py

@@ -20,6 +20,7 @@ class DataType(StrEnum):
     TIMESTAMP = "timestamp"
 
     def to_py(self):
+        """Get python type."""
         match self:
             case DataType.INTEGER:
                 return int
@@ -31,6 +32,7 @@ def to_py(self):
                 return datetime
 
     def to_arrow(self):
+        """Get arrow type."""
         match self:
             case DataType.INTEGER:
                 return pa.int64()
@@ -69,10 +71,35 @@ def __init__(self, path: str | Path):
 
     @property
     def path(self) -> Path:
+        """Cache directory path."""
         return self._path
 
-    @property
-    def files(self) -> list[Path]:
+    @staticmethod
+    def _generate_config_path(path: str | Path) -> Path:
+        """
+        Generate cache configuration path.
+
+        Parameters
+        ----------
+        path : str | Path
+            Where the cache is stored.
+
+        Returns
+        -------
+        Path
+            Configuration filepath.
+        """
+        return Path(path) / ".cfg"
+
+    def get_files(self) -> list[Path]:
+        """
+        Retrieve all files.
+
+        Returns
+        -------
+        list[Path]
+            A list of paths to files in the cache.
+        """
         if not self.path.exists():
             return []
         files = []
@@ -82,26 +109,21 @@ def files(self) -> list[Path]:
                 files.append(Path(full_path))
         return files
 
-    @property
-    def num_files(self) -> int:
-        return len(self.files)
+    def get_dataset_files(self) -> list[Path]:
+        """
+        Retrieve all dataset files.
 
-    @property
-    def dataset_files(self) -> list[Path]:
+        Returns
+        -------
+        list[Path]
+            A list of paths to dataset files in the cache.
+        """
         if not self.path.exists():
             return []
         return [
             Path(filepath) for filepath in glob.glob(f"{self._path}/*.parquet")
         ]
 
-    @property
-    def num_dataset_files(self) -> int:
-        return len(self.dataset_files)
-
-    @staticmethod
-    def _generate_config_path(path: str | Path) -> Path:
-        return Path(path) / ".cfg"
-
 
 class CacheReader(CacheFiles):
     def __init__(
@@ -118,9 +140,17 @@ def __init__(
 
     @classmethod
     def load(cls, path: str | Path):
-        path = Path(path)
+        """
+        Load cache from disk.
+
+        Parameters
+        ----------
+        path : str | Path
+            Where the cache is stored.
+        """
 
         # validate path
+        path = Path(path)
         if not path.exists():
             raise FileNotFoundError(f"Directory does not exist: {path}")
         elif not path.is_dir():
@@ -151,22 +181,27 @@ def _retrieve(config: dict, key: str):
 
     @property
     def dataset(self) -> ds.Dataset:
+        """Cache dataset."""
         return ds.dataset(self._path, format="parquet")
 
     @property
     def schema(self) -> pa.Schema:
+        """Cache schema."""
         return self.dataset.schema
 
     @property
     def batch_size(self) -> int:
+        """Batch size for writes."""
         return self._batch_size
 
     @property
     def rows_per_file(self) -> int:
+        """Target number of rows per file."""
         return self._rows_per_file
 
     @property
     def compression(self) -> str:
+        """File compression method."""
         return self._compression
 
 
@@ -190,15 +225,60 @@ def __init__(
         self._buffer = []
         self._count = 0
 
+    @property
+    def schema(self) -> pa.Schema:
+        """Cache schema."""
+        return self._schema
+
+    @property
+    def dataset(self) -> ds.Dataset:
+        """Cache dataset."""
+        return ds.dataset(
+            self._path,
+            format="parquet",
+            schema=self.schema,
+        )
+
+    @property
+    def batch_size(self) -> int:
+        """Batch size for writes."""
+        return self._batch_size
+
+    @property
+    def rows_per_file(self) -> int:
+        """Target number of rows per file."""
+        return self._rows_per_file
+
+    @property
+    def compression(self) -> str:
+        """File compression method."""
+        return self._compression
+
     @classmethod
     def create(
         cls,
         path: str | Path,
         schema: pa.Schema,
-        batch_size: int = 1000,
-        rows_per_file: int = 10000,
+        batch_size: int = 1_000,
+        rows_per_file: int = 10_000,
         compression: str = "snappy",
     ):
+        """
+        Create a cache.
+
+        Parameters
+        ----------
+        path : str | Path
+            Where to write the cache.
+        schema : pa.Schema
+            Cache schema.
+        batch_size : int, default=1_000
+            Target batch size when writing chunks.
+        rows_per_file : int, default=10_000
+            Target number of rows to store per file.
+        compression : str, default="snappy"
+            Compression method to use when storing on disk.
+        """
         Path(path).mkdir(parents=True, exist_ok=False)
         cfg_path = cls._generate_config_path(path)
         with open(cfg_path, "w") as f:
@@ -218,8 +298,15 @@ def create(
 
     @classmethod
     def load(cls, path: str | Path):
+        """
+        Load cache from disk.
+
+        Parameters
+        ----------
+        path : str | Path
+            Where the cache is stored.
+        """
         path = Path(path)
-        # validate path
         if not path.exists():
             raise FileNotFoundError(f"Directory does not exist: {path}")
         elif not path.is_dir():
@@ -239,6 +326,14 @@ def load(cls, path: str | Path):
 
     @classmethod
     def delete(cls, path: str | Path):
+        """
+        Delete a cache from disk.
+
+        Parameters
+        ----------
+        path : str | Path
+            Where the cache is stored.
+        """
         path = Path(path)
         if not path.exists():
             return
@@ -248,7 +343,7 @@ def delete(cls, path: str | Path):
         if cfg_path.exists() and cfg_path.is_file():
             cfg_path.unlink()
         # delete parquet files
-        for file in cache.dataset_files:
+        for file in cache.get_dataset_files():
             if file.exists() and file.is_file() and file.suffix == ".parquet":
                 file.unlink()
         # delete empty cache directory
@@ -258,6 +353,14 @@ def write_rows(
         self,
         rows: list[dict[str, Any]],
     ):
+        """
+        Write rows to cache.
+
+        Parameters
+        ----------
+        rows : list[dict[str, Any]]
+            A list of rows represented by dictionaries mapping fields to values.
+        """
         if not rows:
             return
         batch = pa.RecordBatch.from_pylist(rows, schema=self.schema)
@@ -267,6 +370,14 @@ def write_batch(
         self,
         batch: pa.RecordBatch | dict[str, list | np.ndarray | pa.Array],
     ):
+        """
+        Write a batch to cache.
+
+        Parameters
+        ----------
+        batch : pa.RecordBatch | dict[str, list | np.ndarray | pa.Array]
+            A batch of columnar data.
+        """
         if isinstance(batch, dict):
             batch = pa.RecordBatch.from_pydict(batch)
 
@@ -303,10 +414,19 @@ def write_table(
         self,
         table: pa.Table,
     ):
+        """
+        Write a table directly to cache.
+
+        Parameters
+        ----------
+        table : pa.Table
+            A populated table.
+        """
         self.flush()
-        pq.write_table(table, where=self._next_filename())
+        pq.write_table(table, where=self._generate_next_filename())
 
     def flush(self):
+        """Flush the cache buffer."""
         if self._buffer:
             combined_arrays = [
                 pa.concat_arrays([b.column(name) for b in self._buffer])
@@ -322,8 +442,9 @@ def flush(self):
         self._count = 0
         self._close_writer()
 
-    def _next_filename(self) -> Path:
-        files = self.dataset_files
+    def _generate_next_filename(self) -> Path:
+        """Generates next dataset filepath."""
+        files = self.get_dataset_files()
         if not files:
             next_index = 0
         else:
@@ -335,7 +456,7 @@ def _get_or_create_writer(self) -> pq.ParquetWriter:
         if self._writer is not None:
             return self._writer
         self._writer = pq.ParquetWriter(
-            where=self._next_filename(),
+            where=self._generate_next_filename(),
             schema=self.schema,
             compression=self.compression,
         )
@@ -354,27 +475,3 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_val, exc_tb):
         """Context manager exit - ensures data is flushed."""
         self.flush()
-
-    @property
-    def schema(self) -> pa.Schema:
-        return self._schema
-
-    @property
-    def dataset(self) -> ds.Dataset:
-        return ds.dataset(
-            self._path,
-            format="parquet",
-            schema=self.schema,
-        )
-
-    @property
-    def batch_size(self) -> int:
-        return self._batch_size
-
-    @property
-    def rows_per_file(self) -> int:
-        return self._rows_per_file
-
-    @property
-    def compression(self) -> str:
-        return self._compression