dayhoff-tools 1.0.0__tar.gz

dayhoff_tools-1.0.0/PKG-INFO

@@ -0,0 +1,122 @@
+Metadata-Version: 2.3
+Name: dayhoff-tools
+Version: 1.0.0
+Summary: Common tools for all the repos at Dayhoff Labs
+Author: Daniel Martin-Alarcon
+Author-email: dma@dayhofflabs.com
+Requires-Python: >=3.10,<4.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Provides-Extra: all
+Provides-Extra: core
+Provides-Extra: dev
+Requires-Dist: biopython (>=1.84) ; extra == "all"
+Requires-Dist: biopython (>=1.84) ; extra == "core"
+Requires-Dist: black (>=25.1.0) ; extra == "all"
+Requires-Dist: black (>=25.1.0) ; extra == "dev"
+Requires-Dist: boto3 (>=1.36.8) ; extra == "all"
+Requires-Dist: boto3 (>=1.36.8) ; extra == "core"
+Requires-Dist: colorlog (>=6.8.2) ; extra == "all"
+Requires-Dist: colorlog (>=6.8.2) ; extra == "dev"
+Requires-Dist: docker (>=7.1.0) ; extra == "all"
+Requires-Dist: docker (>=7.1.0) ; extra == "core"
+Requires-Dist: dvc (>=3.48.2) ; extra == "all"
+Requires-Dist: dvc (>=3.48.2) ; extra == "dev"
+Requires-Dist: dvc-gs (>=3.0.1) ; extra == "all"
+Requires-Dist: dvc-gs (>=3.0.1) ; extra == "dev"
+Requires-Dist: fair-esm (>=2.0.0) ; extra == "all"
+Requires-Dist: fair-esm (>=2.0.0) ; extra == "core"
+Requires-Dist: firebase-admin (>=6.5.0)
+Requires-Dist: flake8 (>=7.0.0) ; extra == "all"
+Requires-Dist: flake8 (>=7.0.0) ; extra == "dev"
+Requires-Dist: h5py (>=3.11.0) ; extra == "all"
+Requires-Dist: h5py (>=3.11.0) ; extra == "core"
+Requires-Dist: isort (>=5.13.2) ; extra == "all"
+Requires-Dist: isort (>=5.13.2) ; extra == "dev"
+Requires-Dist: numpy (<2.0.0) ; extra == "all"
+Requires-Dist: numpy (<2.0.0) ; extra == "dev"
+Requires-Dist: pandas (>=2.2.3) ; extra == "dev"
+Requires-Dist: pylance (>=0.10.2) ; extra == "all"
+Requires-Dist: pylance (>=0.10.2) ; extra == "dev"
+Requires-Dist: pylint (>=3.1.0) ; extra == "all"
+Requires-Dist: pylint (>=3.1.0) ; extra == "dev"
+Requires-Dist: pytest (>=8.0.2) ; extra == "all"
+Requires-Dist: pytest (>=8.0.2) ; extra == "dev"
+Requires-Dist: pytest-cov (>=4.1.0) ; extra == "all"
+Requires-Dist: pytest-cov (>=4.1.0) ; extra == "dev"
+Requires-Dist: pytest-mock (>=3.12.0) ; extra == "all"
+Requires-Dist: pytest-mock (>=3.12.0) ; extra == "dev"
+Requires-Dist: pyyaml (>=6.0)
+Requires-Dist: questionary (>=2.0.1) ; extra == "all"
+Requires-Dist: questionary (>=2.0.1) ; extra == "core"
+Requires-Dist: rdkit-pypi (>=2022.9.5) ; extra == "all"
+Requires-Dist: rdkit-pypi (>=2022.9.5) ; extra == "dev"
+Requires-Dist: requests (>=2.31.0)
+Requires-Dist: torch (>=1.10.0) ; extra == "all"
+Requires-Dist: torch (>=1.10.0) ; extra == "dev"
+Requires-Dist: torchvision (>=0.11.0) ; extra == "all"
+Requires-Dist: torchvision (>=0.11.0) ; extra == "dev"
+Requires-Dist: transformers (>=4.20.0) ; extra == "all"
+Requires-Dist: transformers (>=4.36.2) ; extra == "dev"
+Requires-Dist: typer (>=0.9.0)
+Description-Content-Type: text/markdown
+
+# dayhoff-tools
+
+A set of small, sharp tools for everyone at Dayhoff.
+
+## Hosting and Auth
+
+This repo uses Poetry to build and publish a package to GCP Artifact Registry, at `https://us-central1-python.pkg.dev/enzyme-discovery/pypirate/`.  This depends on a Poetry plugin that's now in the standard chassis setup (`keyrings.google-artifactregistry-auth`), and also on the active service account having read access to Artifact Registry. That much is set up for the standard dev container service account, but may not be available to other intended users.
+
+## CLI commands
+
+Unlike all the repos that use dayhoff-tools, here you have to install the package explicitly before using the CLI:
+
+```sh
+poetry install
+```
+
+## Publish a new version
+
+1. Update version number in `pyproject.toml`
+2. Run `dh wheel`
+3. In other repos, run `poetry update dayhoff-tools`
+
+If you want to overwrite an existing wheel, you'll have to manually delete it from the `dist` folder and also the [Artifact Registry repo](https://console.cloud.google.com/artifacts/python/enzyme-discovery/us-central1/pypirate/dayhoff-tools).
+
+## Install in other repos
+
+Installing this library is tricky because we need GCS authentication and also a couple of plugins to install this with either Pip or Poetry.  These have been incorporated into `chassis`, but it's worth noting here what the various parts are.  All this info came from this [Medium post](https://medium.com/google-cloud/python-packages-via-gcps-artifact-registry-ce1714f8e7c1).
+
+1. Get a Service Account with read access to Artifact Registry (such as `github-actions`, which I made for this purpose).
+2. Export the SA key file, copy it to your repo, and make it available through this envvar: `export GOOGLE_APPLICATION_CREDENTIALS=github_actions_key.json`
+
+### ... with Pip
+
+1. `pip install keyring`
+2. `pip install keyrings.google-artifactregistry-auth`
+3. `pip install --upgrade dayhoff-tools --index-url https://us-central1-python.pkg.dev/enzyme-discovery/pypirate/simple/`
+
+### ... with Poetry
+
+1. Add this plugin: `poetry self add keyrings.google-artifactregistry-auth`
+2. Add these sections to `pyproject.toml`.  Note that dayhoff-tools is in a separate group `pypirate` that installs separately from the others.
+
+   ```toml
+   [tool.poetry.group.pypirate.dependencies]
+   dayhoff-tools = {version = "*", source = "pypirate"}
+
+   [[tool.poetry.source]]
+    name = "pypirate"
+    url = "https://us-central1-python.pkg.dev/enzyme-discovery/pypirate/simple/"
+    priority = "supplemental"
+   ```
+
+3. When building a dev container, or in other circumstances when you can't easily authenticate as above, run `poetry install --without pypirate`.  
+4. Otherwise, just `poetry install`.
+5. To ensure you have the latest version, run `poetry update dayhoff-tools`.
+

dayhoff_tools-1.0.0/README.md

@@ -0,0 +1,55 @@
+# dayhoff-tools
+
+A set of small, sharp tools for everyone at Dayhoff.
+
+## Hosting and Auth
+
+This repo uses Poetry to build and publish a package to GCP Artifact Registry, at `https://us-central1-python.pkg.dev/enzyme-discovery/pypirate/`.  This depends on a Poetry plugin that's now in the standard chassis setup (`keyrings.google-artifactregistry-auth`), and also on the active service account having read access to Artifact Registry. That much is set up for the standard dev container service account, but may not be available to other intended users.
+
+## CLI commands
+
+Unlike all the repos that use dayhoff-tools, here you have to install the package explicitly before using the CLI:
+
+```sh
+poetry install
+```
+
+## Publish a new version
+
+1. Update version number in `pyproject.toml`
+2. Run `dh wheel`
+3. In other repos, run `poetry update dayhoff-tools`
+
+If you want to overwrite an existing wheel, you'll have to manually delete it from the `dist` folder and also the [Artifact Registry repo](https://console.cloud.google.com/artifacts/python/enzyme-discovery/us-central1/pypirate/dayhoff-tools).
+
+## Install in other repos
+
+Installing this library is tricky because we need GCS authentication and also a couple of plugins to install this with either Pip or Poetry.  These have been incorporated into `chassis`, but it's worth noting here what the various parts are.  All this info came from this [Medium post](https://medium.com/google-cloud/python-packages-via-gcps-artifact-registry-ce1714f8e7c1).
+
+1. Get a Service Account with read access to Artifact Registry (such as `github-actions`, which I made for this purpose).
+2. Export the SA key file, copy it to your repo, and make it available through this envvar: `export GOOGLE_APPLICATION_CREDENTIALS=github_actions_key.json`
+
+### ... with Pip
+
+1. `pip install keyring`
+2. `pip install keyrings.google-artifactregistry-auth`
+3. `pip install --upgrade dayhoff-tools --index-url https://us-central1-python.pkg.dev/enzyme-discovery/pypirate/simple/`
+
+### ... with Poetry
+
+1. Add this plugin: `poetry self add keyrings.google-artifactregistry-auth`
+2. Add these sections to `pyproject.toml`.  Note that dayhoff-tools is in a separate group `pypirate` that installs separately from the others.
+
+   ```toml
+   [tool.poetry.group.pypirate.dependencies]
+   dayhoff-tools = {version = "*", source = "pypirate"}
+
+   [[tool.poetry.source]]
+    name = "pypirate"
+    url = "https://us-central1-python.pkg.dev/enzyme-discovery/pypirate/simple/"
+    priority = "supplemental"
+   ```
+
+3. When building a dev container, or in other circumstances when you can't easily authenticate as above, run `poetry install --without pypirate`.  
+4. Otherwise, just `poetry install`.
+5. To ensure you have the latest version, run `poetry update dayhoff-tools`.

dayhoff_tools-1.0.0/dayhoff_tools/chemistry/standardizer.py

@@ -0,0 +1,297 @@
+"""Normalization classes for molecules and reactions."""
+
+from abc import ABC, abstractmethod
+
+from rdkit import Chem, rdBase
+from rdkit.Chem import AllChem
+from rdkit.Chem.MolStandardize import rdMolStandardize  # type: ignore
+
+rdBase.DisableLog("rdApp.error")
+rdBase.DisableLog("rdApp.warning")
+
+
+def is_smiles_aromatic(smiles: str) -> bool:
+    """Check if SMILES string contains aromatic atoms.
+
+    Args:
+        smiles (str): Input SMILES string
+
+    Returns:
+        bool: True if aromatic atoms are found
+    """
+
+    rdmol = Chem.MolFromSmiles(smiles, sanitize=False)  # type: ignore
+    if rdmol is None:
+        raise ValueError("invalid SMILES string")
+    return any(at.GetIsAromatic() for at in rdmol.GetAtoms())
+
+
+class BaseStandardizer(ABC):
+    """Abstract base class  for normalizing molecules and reactions."""
+
+    @abstractmethod
+    def standardize_molecule(self, smiles: str) -> str:
+        """Standardize molecules as SMILES strings.
+
+        Args:
+            smiles (str): Input SMILES string
+
+        Returns:
+            str: Output SMILES string
+        """
+        pass
+
+    def standardize_reaction(self, smiles: str) -> str:
+        """Standardize reactions as SMILES/SMARTS strings.
+
+        Args:
+            smiles (str): Input SMILES/SMARTS string
+
+        Returns:
+            str: Output SMILES/SMARTS string
+        """
+        rdrxn = AllChem.ReactionFromSmarts(smiles, useSmiles=True)  # type: ignore
+        rdrxn1 = AllChem.ChemicalReaction()  # type: ignore
+        for rdmol in rdrxn.GetReactants():
+            smiles1 = Chem.MolToSmiles(rdmol, canonical=False)  # type: ignore
+            smiles2 = self.standardize_molecule(smiles1)
+            rdmol1 = Chem.MolFromSmiles(smiles2, sanitize=False)  # type: ignore
+            rdfrags = Chem.GetMolFrags(rdmol1, asMols=True, sanitizeFrags=False)  # type: ignore
+            if len(rdfrags) == 1:
+                rdrxn1.AddReactantTemplate(rdmol1)
+            else:
+                for rdfrag in rdfrags:
+                    rdrxn1.AddReactantTemplate(rdfrag)
+            # rdrxn1.AddReactantTemplate(rdmol1)
+        for rdmol in rdrxn.GetProducts():
+            smiles1 = Chem.MolToSmiles(rdmol, canonical=False)  # type: ignore
+            smiles2 = self.standardize_molecule(smiles1)
+            rdmol1 = Chem.MolFromSmiles(smiles2, sanitize=False)  # type: ignore
+            rdfrags = Chem.GetMolFrags(rdmol1, asMols=True, sanitizeFrags=False)  # type: ignore
+            if len(rdfrags) == 1:
+                rdrxn1.AddProductTemplate(rdmol1)
+            else:
+                for rdfrag in rdfrags:
+                    rdrxn1.AddProductTemplate(rdfrag)
+        return AllChem.ReactionToSmiles(rdrxn1)  # type: ignore
+
+
+class HypervalentStandardizer(BaseStandardizer):
+    """Standardizer for converting double to single bonds in hypervalent
+    compounds."""
+
+    def standardize_molecule(self, smiles: str) -> str:
+        """Standardize molecules as SMILES strings.
+
+        Args:
+            smiles (str): Input SMILES string
+
+        Returns:
+            str: Output SMILES string
+        """
+        rdmol = Chem.MolFromSmiles(smiles, sanitize=False)  # type: ignore
+        if rdmol is None:
+            raise ValueError(f"Invalid SMILES input '{smiles}'")
+        ret = Chem.SanitizeMol(rdmol, sanitizeOps=Chem.SANITIZE_CLEANUP)  # type: ignore
+        if ret > 0:
+            raise ValueError(f"Sanitization failed for SMILES input '{smiles}'")
+        return Chem.MolToSmiles(rdmol)  # type: ignore
+
+
+class RemoveHsStandardizer(BaseStandardizer):
+    """Standardizer for removing explicit hydrogens from molecules."""
+
+    def standardize_molecule(self, smiles: str) -> str:
+        """Standardize molecules as SMILES strings.
+
+        Args:
+            smiles (str): Input SMILES string
+
+        Returns:
+            str: Output SMILES string
+        """
+        rdmol = Chem.MolFromSmiles(smiles, sanitize=False)  # type: ignore
+        if rdmol is None:
+            raise ValueError(f"Invalid SMILES input '{smiles}'")
+        rdmol1 = Chem.RemoveHs(rdmol, sanitize=False)  # type: ignore
+        ret = Chem.SanitizeMol(rdmol1, sanitizeOps=Chem.SANITIZE_FINDRADICALS)  # type: ignore
+        if ret > 0:
+            raise ValueError(f"Sanitization failed for SMILES input '{smiles}'")
+        return Chem.MolToSmiles(rdmol1, canonical=True)  # type: ignore
+
+
+class KekulizeStandardizer(BaseStandardizer):
+    """Standardizer for kekulizing aromatic compounds."""
+
+    def standardize_molecule(self, smiles: str) -> str:
+        """Standardize molecules as SMILES strings.
+
+        Args:
+            smiles (str): Input SMILES string
+
+        Returns:
+            str: Output SMILES string
+        """
+        rdmol = Chem.MolFromSmiles(smiles, sanitize=False)  # type: ignore
+        if rdmol is None:
+            raise ValueError(f"Invalid SMILES input '{smiles}'")
+        rdmol.UpdatePropertyCache(strict=False)
+        Chem.Kekulize(rdmol, clearAromaticFlags=True)  # type: ignore
+        return Chem.MolToSmiles(rdmol, canonical=True)  # type: ignore
+
+
+class UnchargeStandardizer(BaseStandardizer):
+    """Standardizer for removing charges from molecules by protonation/deprotonation."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._uncharger = rdMolStandardize.Uncharger()
+
+    def standardize_molecule(self, smiles: str) -> str:
+        """Standardize molecules as SMILES strings.
+
+        Args:
+            smiles (str): Input SMILES string
+
+        Returns:
+            str: Output SMILES string
+        """
+        rdmol = Chem.MolFromSmiles(smiles, sanitize=False)  # type: ignore
+        if rdmol is None:
+            raise ValueError(f"Invalid SMILES input '{smiles}'")
+        rdmol1 = self._uncharger.uncharge(rdmol)
+        return Chem.MolToSmiles(rdmol1)  # type: ignore
+
+    def standardize_reaction(self, smiles: str) -> str:
+        """Standardize reactions as SMILES/SMARTS strings.
+
+        Args:
+            smiles (str): Input SMILES/SMARTS string
+
+        Returns:
+            str: Output SMILES/SMARTS string
+        """
+        rdrxn = AllChem.ReactionFromSmarts(smiles, useSmiles=True)  # type: ignore
+        rdrxn1 = AllChem.ChemicalReaction()  # type: ignore
+
+        # Remove all explicit protons from the reaction
+        reactant_total_charge = 0
+        product_total_charge = 0
+        for rdmol in rdrxn.GetReactants():
+            smiles1 = Chem.MolToSmiles(rdmol, canonical=False)  # type: ignore
+            if smiles1 != "[H+]":
+                smiles2 = self.standardize_molecule(smiles1)
+                rdmol1 = Chem.MolFromSmiles(smiles2, sanitize=False)  # type: ignore
+                reactant_total_charge += Chem.GetFormalCharge(rdmol1)  # type: ignore
+                rdrxn1.AddReactantTemplate(rdmol1)
+        for rdmol in rdrxn.GetProducts():
+            smiles1 = Chem.MolToSmiles(rdmol, canonical=False)  # type: ignore
+            if smiles1 != "[H+]":
+                smiles2 = self.standardize_molecule(smiles1)
+                rdmol1 = Chem.MolFromSmiles(smiles2, sanitize=False)  # type: ignore
+                product_total_charge += Chem.GetFormalCharge(rdmol1)  # type: ignore
+                rdrxn1.AddProductTemplate(rdmol1)
+
+        # Rebalance reaction with protons
+        if reactant_total_charge > product_total_charge:
+            rdmol1 = Chem.MolFromSmiles("[H+]", sanitize=False)  # type: ignore
+            for _ in range(reactant_total_charge - product_total_charge):
+                rdrxn1.AddProductTemplate(rdmol1)
+        elif product_total_charge > reactant_total_charge:
+            rdmol1 = Chem.MolFromSmiles("[H+]", sanitize=False)  # type: ignore
+            for _ in range(product_total_charge - reactant_total_charge):
+                rdrxn1.AddReactantTemplate(rdmol1)
+        return AllChem.ReactionToSmiles(rdrxn1)  # type: ignore
+
+
+class MetalStandardizer(BaseStandardizer):
+    """Standardizer for disconnecting bonds between metals and N, O, F atoms."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._disconnector = rdMolStandardize.MetalDisconnector()
+
+    def standardize_molecule(self, smiles: str) -> str:
+        """Standardize molecules as SMILES strings.
+
+        Args:
+            smiles (str): Input SMILES string
+
+        Returns:
+            str: Output SMILES string
+        """
+        rdmol = Chem.MolFromSmiles(smiles, sanitize=False)  # type: ignore
+        if rdmol is None:
+            raise ValueError(f"Invalid SMILES input '{smiles}'")
+
+        flags = Chem.SANITIZE_ALL ^ Chem.SANITIZE_PROPERTIES  # type: ignore
+        if not is_smiles_aromatic(smiles):
+            flags ^= Chem.SANITIZE_SETAROMATICITY  # type: ignore
+        Chem.SanitizeMol(rdmol, sanitizeOps=flags)  # type: ignore
+        rdmol1 = self._disconnector.Disconnect(rdmol)
+        return Chem.MolToSmiles(rdmol1)  # type: ignore
+
+
+class Standardizer(BaseStandardizer):
+    """Aggregate standardizer for molecules and reactions."""
+
+    def __init__(
+        self,
+        *,
+        standardize_hypervalent: bool = True,
+        standardize_remove_hs: bool = True,
+        standardize_kekulize: bool = False,
+        standardize_uncharge: bool = False,
+        standardize_metals: bool = True,
+    ):
+        """Initialize the standardizer.
+
+        Args:
+            standardize_hypervalent (bool): Convert double to single bonds in
+                hypervalent compounds
+            standardize_remove_hs (bool): Remove explicit hydrogen atoms
+            standardize_kekulize (bool): Kekulize aromatic compounds
+            standardize_uncharge (bool): Remove charges from molecules by
+                protonation/deprotonation
+            standardize_metals (bool): Disconnect bonds between metals and
+                N, O, F atoms
+        """
+        self._standardizers = []
+        if standardize_hypervalent:
+            self._standardizers.append(HypervalentStandardizer())
+        if standardize_remove_hs:
+            self._standardizers.append(RemoveHsStandardizer())
+        if standardize_kekulize:
+            self._standardizers.append(KekulizeStandardizer())
+        if standardize_uncharge:
+            self._standardizers.append(UnchargeStandardizer())
+        if standardize_metals:
+            self._standardizers.append(MetalStandardizer())
+
+    def standardize_molecule(self, smiles: str) -> str:
+        """Standardize molecules as SMILES strings.
+
+        Args:
+            smiles (str): Input SMILES string
+
+        Returns:
+            str: Output SMILES string
+        """
+        smiles1 = smiles
+        for standardizer in self._standardizers:
+            smiles1 = standardizer.standardize_molecule(smiles1)
+        return smiles1
+
+    def standardize_reaction(self, smiles: str) -> str:
+        """Standardize reactions as SMILES/SMARTS strings.
+
+        Args:
+            smarts (str): Input SMILES/SMARTS string
+
+        Returns:
+            str: Output SMILES/SMARTS string
+        """
+        smiles1 = smiles
+        for standardizer in self._standardizers:
+            smiles1 = standardizer.standardize_reaction(smiles1)
+        return smiles1

dayhoff_tools-1.0.0/dayhoff_tools/chemistry/utils.py

@@ -0,0 +1,63 @@
+"""Chemistry utils for refinery."""
+
+from typing import Optional
+
+from rdkit import Chem, rdBase
+
+__all__ = ["generate_inchikey"]
+
+rdBase.DisableLog("rdApp.warning")
+rdBase.DisableLog("rdApp.error")
+
+
+def generate_inchikey(
+    s: str, rgroup_smiles: Optional[str] = None, ignore_direction: bool = False
+) -> str:
+    """Generate INChI key from SMILES or reaction SMILES string.
+       Passes exceptions to the caller.
+
+    Args:
+        s (str): SMILES or reaction SMILES
+        rgroup_smiles (Optional[str]): Replacement SMILES string for R groups (*).
+            If None, R groups raise a ValueError.
+        ignore_direction (bool, optional): Ignore direction in reaction SMILES.
+            Has no effect on SMILES. Defaults to False.
+
+    Returns:
+        str: INChI key of molecule (products>>substrates)
+    """
+
+    if ">>" in s:
+        reactants, products = s.split(">>", maxsplit=1)
+        reactants_inchikey = generate_inchikey(reactants, rgroup_smiles=rgroup_smiles)
+        products_inchikey = generate_inchikey(products, rgroup_smiles=rgroup_smiles)
+        if ignore_direction and reactants_inchikey > products_inchikey:
+            reaction_inchikey = products_inchikey + ">>" + reactants_inchikey
+        else:
+            reaction_inchikey = reactants_inchikey + ">>" + products_inchikey
+        return reaction_inchikey
+    elif "*" in s:
+        if rgroup_smiles is not None:
+            replaced_smiles = s.replace("*", rgroup_smiles)
+            if "()" in replaced_smiles:
+                replaced_smiles = replaced_smiles.replace("()", "")
+            return generate_inchikey(replaced_smiles)
+        else:
+            raise ValueError(
+                f"Found R (*) groups in SMILES string {s}. Set rgroup_smiles to replace."
+            )
+    elif s != "":
+        rdmol = None
+        try:
+            rdmol = Chem.MolFromSmiles(s, sanitize=True)  # type: ignore
+        except Exception:
+            pass
+        if rdmol is None:
+            raise ValueError(f"Invalid SMILES string {s}")
+        inchikey = Chem.MolToInchiKey(rdmol)
+        if inchikey != "":
+            return inchikey
+        else:
+            raise ValueError("Could not generate INChI key")
+    else:
+        raise ValueError("Empty SMILES string")

dayhoff_tools-1.0.0/dayhoff_tools/cli/main.py

@@ -0,0 +1,90 @@
+"""Entry file for the CLI, which aggregates and aliases all commands."""
+
+import typer
+from dayhoff_tools.cli.utility_commands import (
+    add_to_warehouse_typer,
+    build_and_upload_wheel,
+    delete_local_branch,
+    get_ancestry,
+    import_from_warehouse_typer,
+    rebuild_devcontainer_file,
+    test_github_actions_locally,
+)
+
+app = typer.Typer()
+
+# Utility commands
+app.command("clean")(delete_local_branch)
+app.command("gha")(test_github_actions_locally)
+app.command("rebuild")(rebuild_devcontainer_file)
+app.command("wadd")(add_to_warehouse_typer)
+app.command("wancestry")(get_ancestry)
+app.command("wheel")(build_and_upload_wheel)
+app.command("wimport")(import_from_warehouse_typer)
+
+
+# Use lazy loading for slow-loading swarm commands
+@app.command("reset")
+def reset_wrapper(
+    firestore_collection: str = typer.Option(prompt=True),
+    old_status: str = typer.Option(default="failed", prompt=True),
+    new_status: str = typer.Option(default="available", prompt=True),
+    delete_old: bool = typer.Option(default=True, prompt=True),
+):
+    """Find all the documents in the database with a given status, and
+    make a new document with the same name and a new status."""
+    from dayhoff_tools.cli.swarm_commands import reset_failed_cards
+
+    reset_failed_cards(firestore_collection, old_status, new_status, delete_old)
+
+
+@app.command("zombie")
+def zombie_wrapper(
+    firestore_collection: str = typer.Option(prompt=True),
+    delete_old: bool = typer.Option(default=True, prompt=True),
+    minutes_threshold: int = typer.Option(default=60, prompt=True),
+):
+    """Find all the documents in the database with status "assigned", and "last_updated"
+    older than a specified threshold, and make a new "available" document for them."""
+    from dayhoff_tools.cli.swarm_commands import reset_zombie_cards
+
+    reset_zombie_cards(firestore_collection, delete_old, minutes_threshold)
+
+
+@app.command("status")
+def status_wrapper(
+    firestore_collection: str = typer.Argument(),
+):
+    """Count the various statuses of items in a given collection."""
+    from dayhoff_tools.cli.swarm_commands import get_firestore_collection_status
+
+    get_firestore_collection_status(firestore_collection)
+
+
+# Deployment commands - use lazy loading but preserve argument passing
+@app.command("deploy")
+def deploy_command(
+    mode: str = typer.Argument(help="Deployment mode. Options: local, shell, batch"),
+    config_path: str = typer.Argument(help="Path to the YAML configuration file"),
+):
+    """Unified deployment command."""
+    from dayhoff_tools.deployment.base import deploy
+
+    deploy(mode, config_path)
+
+
+@app.command("job")
+def run_job_command(
+    mode: str = typer.Argument(
+        default="setup_and_execute",
+        help="Mode to run in: setup (setup only), execute (execute only), or setup_and_execute (both)",
+    )
+):
+    """Run a job."""
+    from dayhoff_tools.deployment.job_runner import run_job
+
+    run_job(mode)
+
+
+if __name__ == "__main__":
+    app()