Python API function to scrape compounds from legacy and v2 stack

docs/source/py_api.md

@@ -1,4 +1,3 @@
-
 # Using the Fragalysis Python API
 
 A work in progress Python package to interface with the Fragalysis web service via the [REST API](api) is available at [github.com/xchem/fragalysis](https://github.com/xchem/fragalysis).
@@ -38,7 +37,7 @@ download_target(name=target_name, tas=target_access_string, token=token, stack="
 The `token` keyword can be ommitted the target is public, `stack` can be either "production", "staging" or the URL of another Fragalysis deployment, tas is the "Target Access String" or DLS proposal-session string (e.g. `lb32627-66`), the destination can be any path and is "." by default.
 
 ```{note}
-The Fragalysis frontend may offer more options for target download than this API. Contributions to update the python API for feature parity are much appreciated. 
+The Fragalysis frontend may offer more options for target download than this API. Contributions to update the python API for feature parity are much appreciated.
 
 The `download_target` method which will need updating is in https://github.com/xchem/fragalysis/blob/main/fragalysis/requests/download.py
 
@@ -82,4 +81,5 @@ Then use the `download_target` function as described above.
 
 fragalysis.requests.download.target_list
 fragalysis.requests.download.download_target
+fragalysis.requests.compounds.get_target_compound_smiles
 ```

fragalysis/requests/__init__.py

@@ -3,3 +3,4 @@
 from .plotly import upload_graph
 from .fragmenstein import fragmenstein_place, fragmenstein_combine
 from .knitwork import knitwork
+from .compounds import get_target_compound_smiles

fragalysis/requests/__main__.py

@@ -0,0 +1,66 @@
+import mrich
+from typer import Typer
+
+app = Typer()
+
+
+@app.command()
+def target_uploads(
+    stack: str = "production",
+    token: str | None = None,
+    statistics_only: bool = False,
+    output: str | None = None,
+) -> None:
+    """
+    Get information for targets uploaded to Fragalysis
+
+    :param stack: Only v2 stacks supported
+    :param token: Optional authentication token
+    :param statistics_only: Return statistics only
+    :param output: Optional path to write pickle data to, if None will print to console
+    """
+
+    from .download import target_uploads
+
+    mrich.h1("target_uploads")
+    mrich.var("stack", stack)
+    mrich.var("token", token)
+    mrich.var("statistics_only", statistics_only)
+    mrich.var("output", output)
+    data = target_uploads(stack=stack, token=token, statistics_only=statistics_only)
+
+    if not output:
+        mrich.print(data)
+    else:
+        import pickle
+
+        mrich.writing(output)
+        pickle.dump(data, open(output, "wb"))
+
+
+@app.command()
+def download_target_uploads(
+    name: str,
+    tas: str,
+    index: int | None = None,
+    stack: str = "production",
+    token: str | None = None,
+    # statistics_only: bool = False,
+    destination: str | None = None,
+):
+
+    from .download import download_target_uploads
+
+    mrich.h1("download_target_uploads")
+    mrich.var("name", name)
+    mrich.var("tas", tas)
+    mrich.var("index", index)
+    mrich.var("stack", stack)
+    mrich.var("token", token)
+    mrich.var("destination", destination)
+
+    download_target_uploads(name=name, tas=tas, index=index, stack=stack, token=token)
+
+
+if __name__ == "__main__":
+    app()

fragalysis/requests/compounds.py

@@ -0,0 +1,91 @@
+from urllib.parse import urljoin
+from .session import _session
+from .urls import COMPOUNDS_URL, TARGETS_URL, SITE_OBSERVATIONS_URL
+import mrich
+
+
+def get_target_compound_smiles(
+    stack: str = "production", token: str | None = None
+) -> dict[str, set[str]]:
+    """
+    Use the /api/target_molecules/ endpoint to get the SMILES of all compounds associated with each target in the legacy stack.
+
+    :param stack: The stack to query. Defaults to "production".
+    :param token: Optional authentication token for the API request. If None, uses default session authentication. Defaults to None.
+    :return: A dictionary mapping target names to sets of associated compound SMILES.
+    :rtype: dict[str, set[str]]
+    """
+
+    with _session(stack=stack, token=token) as session:
+
+        match stack:
+            case "legacy":
+
+                response = session.get(urljoin(session.root, "/api/target_molecules/"))
+
+                if not response.ok:
+                    raise Exception(
+                        f"Failed to get target compounds: {response.status_code} - {response.text}"
+                    )
+
+                data = response.json()
+
+                smiles_by_target = {}
+
+                for target_data in data["results"]:
+
+                    target_name = target_data["title"]
+                    smiles_by_target.setdefault(target_name, set())
+
+                    for molecule_data in target_data["molecules"]:
+                        smiles_by_target[target_name].add(
+                            molecule_data["data"]["smiles"]
+                        )
+
+                return smiles_by_target
+
+            case _:
+
+                # assuming v2 stack
+
+                from .download import target_list
+
+                # targets
+                targets_url = urljoin(session.root, TARGETS_URL)
+                mrich.debug(f"GET {targets_url}")
+                target_response = session.get(targets_url)
+
+                if not target_response.ok:
+                    raise Exception(
+                        "Request failed", targets_url, target_response.status_code
+                    )
+
+                targets_data = target_response.json()
+
+                targets = {t["id"]: t["title"] for t in targets_data["results"]}
+
+                smiles_by_target = {}
+
+                for target_id, target_title in targets.items():
+                    mrich.var(target_title, target_id)
+
+                    smiles_by_target.setdefault(target_title, set())
+
+                    observations_url = urljoin(session.root, SITE_OBSERVATIONS_URL)
+                    mrich.debug(f"GET {observations_url}")
+                    response = session.get(
+                        observations_url, params={"target": target_id}
+                    )
+
+                    if not response.ok:
+                        raise Exception(
+                            "Request failed", observations_url, response.status_code
+                        )
+
+                    data = response.json()
+
+                    for observation in data["results"]:
+                        smiles = observation["smiles"]
+                        smiles_by_target[target_title].add(smiles)
+
+                return smiles_by_target

fragalysis/requests/download.py

@@ -245,7 +245,7 @@ def target_uploads(
     stack: str = "production",
     token: str | None = None,
     statistics_only: bool = False,
-) -> dict[(str,str), list]:
+) -> dict[(str, str), list]:
     """Request a dictionary of uploads keyed by target name and target_access_strings from a Fragalysis deployment
 
     :param stack: shorthand or URL of Fragalysis deployment, defaults to "production"
@@ -280,44 +280,126 @@ def target_uploads(
             formatted.setdefault(key, {})
             formatted[key].setdefault("uploads", [])
 
-            formatted[key]["target_id"]=d["target"]
-            formatted[key]["target_name"]=d["target_name"]
-            formatted[key]["target_access_string"]=d["proposal_number"]
-            formatted[key]["project_id"]=d["project"]
-            
+            formatted[key]["target_id"] = d["target"]
+            formatted[key]["target_name"] = d["target_name"]
+            formatted[key]["target_access_string"] = d["proposal_number"]
+            formatted[key]["project_id"] = d["project"]
+
             # reformat the serialised data
 
-            formatted[key]["uploads"].append(dict(
-                xca_tarball_url=d["tarball"],
-                committer_id=d["committer"],
-                committer_name=d["committer_name"],
-                upload_index=d["upload_version"],
-                data_format=f"{d['data_version_major']}.{d['data_version_minor']}",
-                timestamp=datetime.fromisoformat(d["commit_datetime"].replace("Z", "+00:00")),
-            ))
+            formatted[key]["uploads"].append(
+                dict(
+                    xca_tarball_url=d["tarball"],
+                    committer_id=d["committer"],
+                    committer_name=d["committer_name"],
+                    upload_index=d["upload_version"],
+                    data_format=f"{d['data_version_major']}.{d['data_version_minor']}",
+                    timestamp=datetime.fromisoformat(
+                        d["commit_datetime"].replace("Z", "+00:00")
+                    ),
+                )
+            )
 
     # sort and format the data
 
     for key, d in formatted.items():
 
         new_d = {}
-        
+
         # general information
-        new_d["target_id"]=d["target_id"]
-        new_d["target_name"]=d["target_name"]
-        new_d["target_access_string"]=d["target_access_string"]
-        new_d["project_id"]=d["project_id"]
-        
+        new_d["target_id"] = d["target_id"]
+        new_d["target_name"] = d["target_name"]
+        new_d["target_access_string"] = d["target_access_string"]
+        new_d["project_id"] = d["project_id"]
+
         # sort uploads
         sorted_uploads = sorted(d["uploads"], key=lambda d: d["upload_index"])
 
         # latest statistics
-        new_d["last_upload_index"]=sorted_uploads[-1]["upload_index"]
-        new_d["last_upload_timestamp"]=sorted_uploads[-1]["timestamp"]
+        new_d["last_upload_index"] = sorted_uploads[-1]["upload_index"]
+        new_d["last_upload_timestamp"] = sorted_uploads[-1]["timestamp"]
 
         if not statistics_only:
             new_d["uploads"] = sorted_uploads
-        
+
         formatted[key] = new_d
 
     return formatted
+
+
+def download_target_uploads(
+    name: str,
+    tas: str,
+    index: int | None = None,
+    stack: str = "production",
+    token: str | None = None,
+    destination: str = ".",
+):
+    import tarfile
+
+    destination = Path(destination)
+    assert destination.exists()
+    destination = destination / f"{name}_{tas}"
+
+    mrich.var("destination", destination)
+
+    all_uploads = target_uploads(stack=stack, token=token)
+
+    uploads = all_uploads.get((name, tas))
+
+    if not uploads:
+        mrich.error(f"No uploads for {name=} {tas=}")
+        return
+
+    uploads = uploads["uploads"]
+
+    mrich.var("#uploads", len(uploads))
+
+    if index:
+        uploads = [d for d in uploads if d["upload_index"] == index]
+
+    with _session(stack=stack, token=token) as session:
+        for upload in uploads:
+
+            mrich.print(upload)
+
+            # dump the tarball
+
+            tarball_url = upload["xca_tarball_url"]
+            version_out = destination / ("v" + upload["data_format"])
+            tarball_out = version_out / Path(tarball_url).name
+
+            if not (parent := tarball_out.parent).exists():
+                mrich.writing(parent)
+                parent.mkdir(parents=True)
+
+            if not tarball_out.exists():
+
+                mrich.writing(tarball_out)
+
+                with session.get(
+                    tarball_url,
+                    stream=True,
+                ) as r:
+
+                    r.raise_for_status()
+                    with open(tarball_out, "wb") as f:
+                        for i, chunk in mrich.track(
+                            enumerate(r.iter_content(chunk_size=8192)),
+                            prefix="Downloading",
+                        ):
+                            if i % 100 == 0:
+                                mrich.set_progress_field("chunks", i)
+                            f.write(chunk)
+
+            with mrich.loading("Expanding tarball..."):
+
+                extract_dir = version_out / f"upload_{upload['upload_index']}"
+
+                if extract_dir.exists():
+                    mrich.warning(f"Skipping existing {extract_dir=}")
+                    continue
+
+                mrich.writing(extract_dir)
+                with tarfile.open(tarball_out, "r:gz") as tar_ref:
+                    tar_ref.extractall(version_out)

fragalysis/requests/urls.py

@@ -4,12 +4,14 @@
     "staging": "https://fragalysis.xchem.diamond.ac.uk",
     "production": "https://fragalysis.diamond.ac.uk",
     "matej-dev": "https://fragalysis-matej-default.xchem-dev.diamond.ac.uk",
+    "legacy": "https://fragalysis-legacy.xchem.diamond.ac.uk",
 }
 
 LOGIN_URL = "/accounts/login/"
 DOWNLOAD_URL = "/api/download_structures/"
 LANDING_PAGE_URL = "/viewer/react/landing/"
 TARGETS_URL = "/api/targets/"
+COMPOUNDS_URL = "/api/compounds/"
 PROJECTS_URL = "/api/projects/"
 SESSION_PROJECTS_URL = "/api/session-projects/"
 SNAPSHOTS_URL = "/api/snapshots/"

setup.py

@@ -1,11 +1,15 @@
 from setuptools import setup
 
 setup(
-   name='fragalysis',
-   version='1.0',
-   description='Python module to interact with Fragalysis',
-   author='Diamond Light Source / Informatics Matters',
-   author_email='max.winokan@diamond.ac.uk',
-   packages=['fragalysis'],  #same as name
-   install_requires=['ipywidgets', 'mpytools'], #external packages as dependencies
-)
+    name="fragalysis",
+    version="1.0",
+    description="Python module to interact with Fragalysis",
+    author="Diamond Light Source / Informatics Matters",
+    author_email="max.winokan@diamond.ac.uk",
+    packages=["fragalysis"],  # same as name
+    install_requires=[
+        "ipywidgets",
+        "mrich",
+        "typer",
+    ],  # external packages as dependencies
+)