fix: handle missing terms-of-use columns in manifest (#766)

Author: adilraza99

malariagen_data/anoph/base.py

@@ -562,6 +562,10 @@ def _sample_set_has_unrestricted_use(self, *, sample_set: str):
         release_manifest_df = self._read_sample_sets_manifest(
             single_release=sample_set_release
         )
+
+        if "unrestricted_use" not in release_manifest_df.columns:
+            return False
+
         sample_set_records_srs = release_manifest_df.loc[
             release_manifest_df["sample_set"] == sample_set, "unrestricted_use"
         ]
@@ -824,13 +828,27 @@ def lookup_study_info(self, sample_set: base_params.sample_set) -> dict:
     def lookup_terms_of_use_info(self, sample_set: base_params.sample_set) -> dict:
         if self._cache_sample_set_to_terms_of_use_info is None:
             df_sample_sets = self._available_sample_sets().set_index("sample_set")
-            self._cache_sample_set_to_terms_of_use_info = df_sample_sets[
-                [
-                    "terms_of_use_expiry_date",
-                    "terms_of_use_url",
-                    "unrestricted_use",
-                ]
-            ].to_dict(orient="index")
+            expected_cols = [
+                "terms_of_use_expiry_date",
+                "terms_of_use_url",
+                "unrestricted_use",
+            ]
+            placeholder_values = {
+                "terms_of_use_expiry_date": "2099-12-31",
+                "terms_of_use_url": float("nan"),
+                "unrestricted_use": False,
+            }
+            available_cols = [c for c in expected_cols if c in df_sample_sets.columns]
+            if available_cols:
+                lookup = df_sample_sets[available_cols].to_dict(orient="index")
+                missing_cols = set(expected_cols) - set(available_cols)
+                if missing_cols:
+                    for ss in lookup:
+                        for mc in missing_cols:
+                            lookup[ss][mc] = placeholder_values[mc]
+            else:
+                lookup = {ss: dict(placeholder_values) for ss in df_sample_sets.index}
+            self._cache_sample_set_to_terms_of_use_info = lookup
         try:
             return self._cache_sample_set_to_terms_of_use_info[sample_set]
         except KeyError as e:

tests/anoph/test_base.py

@@ -256,3 +256,119 @@ def test_lookup_study(fixture, api):
 
     with pytest.raises(ValueError):
         api.lookup_study("foobar")
+
+
+def _strip_terms_of_use_from_manifest(manifest_path):
+    """Rewrite a manifest TSV file without terms-of-use columns."""
+    df = pd.read_csv(manifest_path, sep="\t")
+    cols_to_drop = [c for c in df.columns if c.startswith("terms_of_use")]
+    df = df.drop(columns=cols_to_drop)
+    df.to_csv(manifest_path, index=False, sep="\t")
+
+
+def test_lookup_terms_of_use_info_missing_columns(ag3_sim_fixture):
+    import shutil
+
+    manifest_paths = [
+        ag3_sim_fixture.bucket_path / "v3" / "manifest.tsv",
+        ag3_sim_fixture.bucket_path / "v3.1" / "manifest.tsv",
+    ]
+    backups = []
+    for mp in manifest_paths:
+        bp = mp.parent / "manifest.tsv.bak"
+        shutil.copy2(mp, bp)
+        backups.append(bp)
+
+    try:
+        for mp in manifest_paths:
+            _strip_terms_of_use_from_manifest(mp)
+
+        api = AnophelesBase(
+            url=ag3_sim_fixture.url,
+            public_url=ag3_sim_fixture.url,
+            config_path=_ag3.CONFIG_PATH,
+            major_version_number=_ag3.MAJOR_VERSION_NUMBER,
+            major_version_path=_ag3.MAJOR_VERSION_PATH,
+            pre=True,
+        )
+
+        sample_set = "1177-VO-ML-LEHMANN-VMF00004"
+        info = api.lookup_terms_of_use_info(sample_set)
+        assert isinstance(info, dict)
+        assert "terms_of_use_expiry_date" in info
+        assert "terms_of_use_url" in info
+        assert "unrestricted_use" in info
+        assert info["terms_of_use_expiry_date"] == "2099-12-31"
+        assert pd.isna(info["terms_of_use_url"])
+        assert info["unrestricted_use"] is False
+    finally:
+        for mp, bp in zip(manifest_paths, backups):
+            shutil.move(bp, mp)
+
+
+def test_sample_set_has_unrestricted_use_missing_column(ag3_sim_fixture):
+    import shutil
+
+    manifest_paths = [
+        ag3_sim_fixture.bucket_path / "v3" / "manifest.tsv",
+        ag3_sim_fixture.bucket_path / "v3.1" / "manifest.tsv",
+    ]
+    backups = []
+    for mp in manifest_paths:
+        bp = mp.parent / "manifest.tsv.bak"
+        shutil.copy2(mp, bp)
+        backups.append(bp)
+
+    try:
+        for mp in manifest_paths:
+            _strip_terms_of_use_from_manifest(mp)
+
+        api = AnophelesBase(
+            url=ag3_sim_fixture.url,
+            public_url=ag3_sim_fixture.url,
+            config_path=_ag3.CONFIG_PATH,
+            major_version_number=_ag3.MAJOR_VERSION_NUMBER,
+            major_version_path=_ag3.MAJOR_VERSION_PATH,
+            pre=True,
+        )
+
+        sample_set = "1177-VO-ML-LEHMANN-VMF00004"
+        result = api._sample_set_has_unrestricted_use(sample_set=sample_set)
+        assert result is False
+    finally:
+        for mp, bp in zip(manifest_paths, backups):
+            shutil.move(bp, mp)
+
+
+def test_sample_sets_no_terms_of_use(ag3_sim_fixture):
+    import shutil
+
+    manifest_paths = [
+        ag3_sim_fixture.bucket_path / "v3" / "manifest.tsv",
+        ag3_sim_fixture.bucket_path / "v3.1" / "manifest.tsv",
+    ]
+    backups = []
+    for mp in manifest_paths:
+        bp = mp.parent / "manifest.tsv.bak"
+        shutil.copy2(mp, bp)
+        backups.append(bp)
+
+    try:
+        for mp in manifest_paths:
+            _strip_terms_of_use_from_manifest(mp)
+
+        api = AnophelesBase(
+            url=ag3_sim_fixture.url,
+            public_url=ag3_sim_fixture.url,
+            config_path=_ag3.CONFIG_PATH,
+            major_version_number=_ag3.MAJOR_VERSION_NUMBER,
+            major_version_path=_ag3.MAJOR_VERSION_PATH,
+            pre=True,
+        )
+
+        df = api.sample_sets(release="3.1")
+        assert isinstance(df, pd.DataFrame)
+        assert len(df) > 0
+    finally:
+        for mp, bp in zip(manifest_paths, backups):
+            shutil.move(bp, mp)