Merge pull request #713 from xchem/m2ms-1633-download-scripts

Add scripts to LHS downloads (issue 1633)

viewer/download_structures.py

@@ -18,6 +18,7 @@
 from typing import Any, Dict
 
 import pandoc
+import requests
 from django.conf import settings
 from django.db.models import Exists, F, OuterRef, Value
 from django.db.models.fields import CharField
@@ -29,6 +30,9 @@
 from .serializers import DownloadStructuresSerializer
 from .tags import get_metadata_fields
 
+# from urllib.parse import urlsplit
+
+
 logger = logging.getLogger(__name__)
 
 # Length of time to keep records of dynamic links.
@@ -63,6 +67,11 @@
     'readme': (''),
 }
 
+# {urls:path in archive} to scripts to be automatically included in downloads
+_SCRIPTS = {
+    'https://github.com/xchem/fragalysis-pymol-scripts': 'pymol',
+}
+
 
 @dataclass(frozen=True)
 class ArchiveFile:
@@ -111,6 +120,33 @@ class ArchiveFile:
 _METADATA_FILE = 'metadata.csv'
 
 
+def _additional_scripts_zip(ziparchive, scripts) -> None:
+    for script_url, script_path in scripts.items():
+        zip_url = script_url.rstrip('/') + '/archive/refs/heads/main.zip'
+        response = requests.get(zip_url)
+        try:
+            response.raise_for_status()
+        except Exception as exc:
+            # repo_name = Path(urlsplit(script_url).path).name
+            ziparchive.writestr(f'scripts/ERROR_DOWNLOADING_{script_path}', str(exc))
+
+        try:
+            with zipfile.ZipFile(BytesIO(response.content)) as zip_file:
+                script_path = Path('scripts', script_path)
+                for zip_info in zip_file.infolist():
+                    if zip_info.is_dir():
+                        continue  # Skip directories
+                    file_data = zip_file.read(zip_info.filename)
+                    archive_path = script_path.joinpath(
+                        *Path(zip_info.filename).parts[1:]
+                    )
+                    # ziparchive.writestr(f'scripts/{zip_info.filename}', file_data)
+                    ziparchive.writestr(str(archive_path), file_data)
+        except zipfile.BadZipFile as exc:
+            # repo_name = Path(urlsplit(script_url).path).name
+            ziparchive.writestr(f'scripts/ERROR_DOWNLOADING_{script_path}', str(exc))
+
+
 def _is_mol_or_sdf(path):
     """Returns True if the file and path look like a MOL or SDF file.
     It does this by simply checking the file's extension.
@@ -646,6 +682,8 @@ def _create_structures_zip(
 
         _document_file_zip(ziparchive, download_path, original_search, host)
 
+        _additional_scripts_zip(ziparchive, _SCRIPTS)
+
         error_file.close()
         if errors > 0:
             logger.warning('errors=%s Adding %s to ziparchive', errors, _ERROR_FILE)