Lale/release v2.4.0

cosmotech/coal/__init__.py

@@ -5,4 +5,4 @@
 # etc., to any person is prohibited unless it has been previously and
 # specifically authorized by written means by Cosmo Tech.
 
-__version__ = "2.3.1"
+__version__ = "2.4.0"

cosmotech/coal/cosmotech_api/apis/dataset.py

@@ -35,6 +35,14 @@ def __init__(
         LOGGER.debug(T("coal.cosmotech_api.initialization.dataset_api_initialized"))
 
     def download_dataset(self, dataset_id) -> Dataset:
+        """Download all parts of a dataset to the configured dataset folder.
+
+        Args:
+            dataset_id: The ID of the dataset to download
+
+        Returns:
+            The Dataset object as returned by the API
+        """
         LOGGER.debug(f"Downloading dataset {dataset_id}")
         dataset = self.get_dataset(
             organization_id=self.configuration.cosmotech.organization_id,
@@ -48,6 +56,14 @@ def download_dataset(self, dataset_id) -> Dataset:
         return dataset
 
     def download_parameter(self, dataset_id) -> Dataset:
+        """Download all parts of a dataset to the configured parameters folder.
+
+        Args:
+            dataset_id: The ID of the dataset to download
+
+        Returns:
+            The Dataset object as returned by the API
+        """
         LOGGER.debug(f"Downloading dataset {dataset_id}")
         dataset = self.get_dataset(
             organization_id=self.configuration.cosmotech.organization_id,
@@ -80,6 +96,15 @@ def _download_part(self, dataset_id, dataset_part, destination):
 
     @staticmethod
     def path_to_parts(_path, part_type) -> list[tuple[str, Path, DatasetPartTypeEnum]]:
+        """Turn a file or directory path into a list of dataset part descriptors.
+
+        Args:
+            _path: Path to a file, or to a directory whose files are each turned into a part
+            part_type: The DatasetPartTypeEnum to assign to every resulting part
+
+        Returns:
+            A list of tuples of (relative part name, absolute file path, part type)
+        """
         if (_path := Path(_path)).is_dir():
             return list((str(_p.relative_to(_path)), _p, part_type) for _p in _path.rglob("*") if _p.is_file())
         return list(((_path.name, _path, part_type),))

cosmotech/coal/cosmotech_api/apis/run.py

@@ -34,5 +34,18 @@ def get_run_metadata(
         include: Optional[list[str]] = None,
         exclude: Optional[list[str]] = None,
     ) -> dict[str, Any]:
+        """Fetch a run and serialize it to a JSON-compatible dictionary.
+
+        Args:
+            organization_id: The ID of the organization
+            workspace_id: The ID of the workspace
+            runner_id: The ID of the runner
+            run_id: The ID of the run
+            include: Optional list of field names to include in the result
+            exclude: Optional list of field names to exclude from the result
+
+        Returns:
+            A dictionary representation of the run's metadata
+        """
         run = self.get_run(organization_id, workspace_id, runner_id, run_id)
         return run.model_dump(by_alias=True, exclude_none=True, include=include, exclude=exclude, mode="json")

cosmotech/coal/cosmotech_api/apis/runner.py

@@ -33,6 +33,16 @@ def get_runner_metadata(
         include: Optional[list[str]] = None,
         exclude: Optional[list[str]] = None,
     ) -> dict[str, Any]:
+        """Fetch a runner and serialize it to a JSON-compatible dictionary.
+
+        Args:
+            runner_id: The ID of the runner, defaults to the configured runner ID
+            include: Optional list of field names to include in the result
+            exclude: Optional list of field names to exclude from the result
+
+        Returns:
+            A dictionary representation of the runner's metadata
+        """
         runner = self.get_runner(
             self.configuration.cosmotech.organization_id,
             self.configuration.cosmotech.workspace_id,
@@ -45,6 +55,11 @@ def download_runner_data(
         self,
         download_datasets: Optional[str] = None,
     ):
+        """Download a runner's parameters and, optionally, its associated datasets.
+
+        Args:
+            download_datasets: If set, also download the runner's base datasets
+        """
         LOGGER.info(T("coal.cosmotech_api.runner.starting_download"))
 
         # Get runner data

cosmotech/coal/cosmotech_api/apis/workspace.py

@@ -32,6 +32,19 @@ def list_filtered_workspace_files(
         workspace_id: str,
         file_prefix: str,
     ) -> list[str]:
+        """List workspace files whose name starts with the given prefix.
+
+        Args:
+            organization_id: The ID of the organization
+            workspace_id: The ID of the workspace
+            file_prefix: The prefix to filter workspace file names by
+
+        Returns:
+            List of matching workspace file names
+
+        Raises:
+            ValueError: If no workspace file matches the given prefix
+        """
         target_list = []
         LOGGER.info(T("coal.cosmotech_api.workspace.target_is_folder"))
         wsf = self.list_workspace_files(organization_id, workspace_id)
@@ -60,6 +73,20 @@ def download_workspace_file(
         file_name: str,
         target_dir: Path,
     ) -> Path:
+        """Download a single workspace file to a local directory.
+
+        Args:
+            organization_id: The ID of the organization
+            workspace_id: The ID of the workspace
+            file_name: The name of the workspace file to download
+            target_dir: The local directory to download the file into
+
+        Returns:
+            The local path of the downloaded file
+
+        Raises:
+            ValueError: If target_dir is not a directory
+        """
         if target_dir.is_file():
             raise ValueError(T("coal.common.file_operations.not_directory").format(target_dir=target_dir))
 
@@ -85,6 +112,22 @@ def upload_workspace_file(
         workspace_path: str,
         overwrite: bool = True,
     ) -> str:
+        """Upload a local file to a workspace.
+
+        Args:
+            organization_id: The ID of the organization
+            workspace_id: The ID of the workspace
+            file_path: Local path of the file to upload
+            workspace_path: Destination path (or directory ending with '/') in the workspace
+            overwrite: If True, overwrite an existing file at the destination
+
+        Returns:
+            The name of the uploaded workspace file
+
+        Raises:
+            ValueError: If file_path does not exist or is not a single file
+            ApiException: If the API call fails, e.g. because the file already exists
+        """
         target_file = Path(file_path)
         if not target_file.exists():
             LOGGER.error(T("coal.common.file_operations.not_exists").format(file_path=file_path))

cosmotech/coal/cosmotech_api/objects/connection.py

@@ -53,6 +53,17 @@ def __init__(
         self.api_client, self.api_type = self.get_api_client()
 
     def get_api_client(self) -> (cosmotech_api.ApiClient, str):
+        """Build an authenticated API client from available environment variables.
+
+        Detects Keycloak, Cosmo Tech API key, or Azure Entra credentials from the
+        environment and builds the corresponding API client.
+
+        Returns:
+            A tuple of the configured ApiClient and a label describing the connection type used
+
+        Raises:
+            EnvironmentError: If no complete set of required environment variables is found
+        """
         existing_keys = set(os.environ.keys())
         missing_azure_keys = self.__azure_env_keys - existing_keys
         missing_api_keys = self.__api_env_keys - existing_keys

cosmotech/coal/cosmotech_api/objects/parameters.py

@@ -84,6 +84,14 @@ def write_parameters_to_json(
         self,
         parameter_folder: str,
     ) -> str:
+        """Write the parameters list to a parameters.json file.
+
+        Args:
+            parameter_folder: Directory in which to create the JSON file
+
+        Returns:
+            The path of the written JSON file
+        """
         pathlib.Path(parameter_folder).mkdir(exist_ok=True, parents=True)
         tmp_parameter_file = os.path.join(parameter_folder, "parameters.json")
 
@@ -98,6 +106,14 @@ def write_parameters_to_csv(
         self,
         parameter_folder: str,
     ) -> str:
+        """Write the parameters list to a parameters.csv file.
+
+        Args:
+            parameter_folder: Directory in which to create the CSV file
+
+        Returns:
+            The path of the written CSV file
+        """
         pathlib.Path(parameter_folder).mkdir(exist_ok=True, parents=True)
         tmp_parameter_file = os.path.join(parameter_folder, "parameters.csv")
 
@@ -116,6 +132,16 @@ def write_parameters(
         write_csv: bool = True,
         write_json: bool = False,
     ) -> Dict[str, str]:
+        """Write the parameters list to CSV and/or JSON files.
+
+        Args:
+            parameter_folder: Directory in which to create the files
+            write_csv: If True, write a parameters.csv file
+            write_json: If True, write a parameters.json file
+
+        Returns:
+            Dictionary mapping format name ("csv", "json") to the written file path
+        """
         result = {}
 
         if write_csv:

cosmotech/coal/store/output/aws_channel.py

@@ -15,6 +15,8 @@
 
 
 class AwsChannel(ChannelInterface):
+    """Output channel that sends the store's tables or database file to an S3 bucket."""
+
     required_keys = {
         "coal": ["store"],
         "cosmotech": ["runner_id"],
@@ -28,7 +30,17 @@ def __init__(self, dct: Dotdict = None):
         self._s3.file_prefix = self.configuration.cosmotech.runner_id + "/" + self._s3.file_prefix
 
     def send(self, filter: Optional[list[str]] = None) -> bool:
+        """Upload the store's database file or tables to the configured S3 bucket.
+
+        Args:
+            filter: Optional list of table names to restrict the upload to
+
+        Returns:
+            True if the data was uploaded successfully
 
+        Raises:
+            ValueError: If the configured output type is not sqlite, csv or parquet
+        """
         _s = Store(configuration=self.configuration)
 
         if self._s3.output_type not in ("sqlite", "csv", "parquet"):
@@ -67,4 +79,5 @@ def send(self, filter: Optional[list[str]] = None) -> bool:
                 )
 
     def delete(self):
+        """Delete all objects previously uploaded to the configured S3 bucket."""
         self._s3.delete_objects()

cosmotech/coal/store/output/az_storage_channel.py

@@ -8,6 +8,8 @@
 
 
 class AzureStorageChannel(ChannelInterface):
+    """Output channel that sends the store's tables to an Azure Blob Storage container."""
+
     required_keys = {
         "coal": ["store"],
         "cosmotech": ["runner_id"],
@@ -31,10 +33,19 @@ def __init__(self, dct: Dotdict = None):
             self.configuration.azure.file_prefix = runner_id + "/" + prefix
 
     def send(self, filter: Optional[list[str]] = None) -> bool:
+        """Dump the store's tables to the configured Azure Blob Storage container.
+
+        Args:
+            filter: Optional list of table names to restrict the dump to
+
+        Returns:
+            True if the data was uploaded successfully
+        """
         dump_store_to_azure(
             self.configuration,
             selected_tables=filter,
         )
 
     def delete(self):
+        """Delete all blobs previously uploaded to the configured Azure Blob Storage container."""
         delete_azure_blobs(self.configuration)

cosmotech/coal/store/output/channel_interface.py

@@ -6,6 +6,8 @@
 
 
 class ChannelInterface:
+    """Base class for an output channel that can send the store's data to a target and delete it."""
+
     required_keys = {}
     requirement_string: str = T("coal.store.output.data_interface.requirements")
 
@@ -21,6 +23,11 @@ def delete(self) -> bool:
         raise NotImplementedError()
 
     def is_available(self) -> bool:
+        """Check whether all required configuration keys are present, falling back to global config.
+
+        Returns:
+            True if the channel is fully configured and can be used
+        """
         try:
             for section in self.required_keys.keys():
                 for key in self.required_keys[section]:

cosmotech/coal/store/output/channel_spliter.py

@@ -11,6 +11,8 @@
 
 
 class ChannelSpliter(ChannelInterface):
+    """Channel that fans out the store's data to every configured output channel."""
+
     requirement_string: str = "(Requires any working interface)"
     targets = list()
     available_interfaces: dict[str, ChannelInterface] = {
@@ -39,6 +41,14 @@ def __init__(self, dct: Dotdict = None):
             raise AttributeError(T("coal.store.output.split.no_targets"))
 
     def send(self, filter: Optional[list[str]] = None) -> bool:
+        """Send the store's data through every available target channel.
+
+        Args:
+            filter: Optional list of table names to restrict the send to
+
+        Returns:
+            True if at least one target channel sent its data successfully
+        """
         any_ok = False
         for i in self.targets:
             try:
@@ -50,6 +60,11 @@ def send(self, filter: Optional[list[str]] = None) -> bool:
         return any_ok
 
     def delete(self) -> bool:
+        """Delete the data previously sent through every available target channel.
+
+        Returns:
+            True if at least one target channel deleted its data successfully
+        """
         any_ok = False
         for i in self.targets:
             try:

cosmotech/coal/store/output/postgres_channel.py

@@ -14,6 +14,8 @@
 
 
 class PostgresChannel(ChannelInterface):
+    """Output channel that sends the store's tables and runner metadata to a PostgreSQL database."""
+
     required_keys = {
         "coal": ["store"],
         "cosmotech": ["organization_id", "workspace_id", "runner_id"],
@@ -34,6 +36,14 @@ def __init__(self, dct: Dotdict = None):
         self.configuration.setup_db = self.configuration.safe_get("setup_db", True)
 
     def send(self, filter: Optional[list[str]] = None) -> bool:
+        """Send the runner's metadata and the store's tables to the configured PostgreSQL database.
+
+        Args:
+            filter: Optional list of table names to restrict the send to
+
+        Returns:
+            True if the data was sent successfully
+        """
         if self.configuration.setup_db:
             create_metadata(self.configuration)
 
@@ -55,5 +65,6 @@ def send(self, filter: Optional[list[str]] = None) -> bool:
             add_fk_constraints(self.configuration)
 
     def delete(self):
+        """Delete the runner's metadata, cascading to delete its associated data in PostgreSQL."""
         # removing metadata will trigger cascade delete on real data
         remove_runner_metadata_from_postgresql(self.configuration)