add module tests/test_package.py with test_egg_sources()

checks that we're correctly packaging data files

Author: janosh

aviary/__init__.py

@@ -0,0 +1,4 @@
+from os.path import abspath, dirname
+
+PKG_DIR = dirname(abspath(__file__))
+ROOT = dirname(PKG_DIR)

aviary/cgcnn/data.py

@@ -3,8 +3,8 @@
 import ast
 import functools
 import json
+import os
 from itertools import groupby
-from os.path import abspath, dirname, exists, join
 from typing import Any, Sequence
 
 import numpy as np
@@ -14,6 +14,8 @@
 from torch import LongTensor, Tensor
 from torch.utils.data import Dataset
 
+from aviary import PKG_DIR
+
 
 class CrystalGraphData(Dataset):
     """Dataset class for the CGCNN structure model."""
@@ -60,12 +62,9 @@ def __init__(
         self.max_num_nbr = max_num_nbr
 
         if elem_emb in ["matscholar200", "cgcnn92", "megnet16", "onehot112"]:
-            elem_emb = join(
-                dirname(abspath(__file__)), f"../embeddings/element/{elem_emb}.json"
-            )
-        else:
-            if not exists(elem_emb):
-                raise AssertionError(f"{elem_emb} does not exist!")
+            elem_emb = f"{PKG_DIR}/embeddings/element/{elem_emb}.json"
+        elif not os.path.exists(elem_emb):
+            raise AssertionError(f"{elem_emb} does not exist!")
 
         with open(elem_emb) as f:
             self.elem_features = json.load(f)

aviary/roost/data.py

@@ -2,7 +2,7 @@
 
 import functools
 import json
-from os.path import abspath, dirname, exists, join
+import os
 from typing import Any, Sequence
 
 import numpy as np
@@ -12,6 +12,8 @@
 from torch import LongTensor, Tensor
 from torch.utils.data import Dataset
 
+from aviary import PKG_DIR
+
 
 class CompositionData(Dataset):
     """Dataset class for the Roost composition model."""
@@ -46,12 +48,9 @@ def __init__(
         self.df = df
 
         if elem_emb in ["matscholar200", "cgcnn92", "megnet16", "onehot112"]:
-            elem_emb = join(
-                dirname(abspath(__file__)), f"../embeddings/element/{elem_emb}.json"
-            )
-        else:
-            if not exists(elem_emb):
-                raise AssertionError(f"{elem_emb} does not exist!")
+            elem_emb = f"{PKG_DIR}/embeddings/element/{elem_emb}.json"
+        elif not os.path.exists(elem_emb):
+            raise AssertionError(f"{elem_emb} does not exist!")
 
         with open(elem_emb) as f:
             self.elem_features = json.load(f)

aviary/wren/data.py

@@ -2,9 +2,9 @@
 
 import functools
 import json
+import os
 import re
 from itertools import groupby
-from os.path import abspath, dirname, exists, join
 from typing import Any, Sequence
 
 import numpy as np
@@ -13,6 +13,7 @@
 from torch import LongTensor, Tensor
 from torch.utils.data import Dataset
 
+from aviary import PKG_DIR
 from aviary.wren.utils import mult_dict, relab_dict
 
 
@@ -51,25 +52,19 @@ def __init__(
         self.df = df
 
         if elem_emb in ["matscholar200", "cgcnn92", "megnet16", "onehot112"]:
-            elem_emb = join(
-                dirname(abspath(__file__)), f"../embeddings/element/{elem_emb}.json"
-            )
-        else:
-            if not exists(elem_emb):
-                raise AssertionError(f"{elem_emb} does not exist!")
+            elem_emb = f"{PKG_DIR}/embeddings/element/{elem_emb}.json"
+        elif not os.path.exists(elem_emb):
+            raise AssertionError(f"{elem_emb} does not exist!")
 
         with open(elem_emb) as f:
             self.elem_features = json.load(f)
 
         self.elem_emb_len = len(list(self.elem_features.values())[0])
 
         if sym_emb in ["bra-alg-off", "spg-alg-off"]:
-            sym_emb = join(
-                dirname(abspath(__file__)), f"../embeddings/wyckoff/{sym_emb}.json"
-            )
-        else:
-            if not exists(sym_emb):
-                raise AssertionError(f"{sym_emb} does not exist!")
+            sym_emb = f"{PKG_DIR}/embeddings/wyckoff/{sym_emb}.json"
+        elif not os.path.exists(sym_emb):
+            raise AssertionError(f"{sym_emb} does not exist!")
 
         with open(sym_emb) as f:
             self.sym_features = json.load(f)

setup.py

@@ -6,14 +6,16 @@
     author="Rhys Goodall",
     author_email="[email protected]",
     url="https://github.com/CompRhys/aviary",
-    description="A Collection of Machine Learning Models for Materials Discovery",
+    description="A collection of machine learning models for materials discovery",
     long_description=open("README.md").read(),
     long_description_content_type="text/markdown",
     packages=find_packages(include=["aviary*"]),
     classifiers=[
         "Programming Language :: Python :: 3.7",
         "Programming Language :: Python :: 3.8",
         "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
         "License :: OSI Approved :: MIT License",
         "Operating System :: OS Independent",
     ],

tests/test_package.py

@@ -0,0 +1,24 @@
+import os
+from glob import glob
+
+import pytest
+
+from aviary import ROOT
+
+package_sources_path = f"{ROOT}/aviary.egg-info/SOURCES.txt"
+
+
+@pytest.mark.skipif(
+    not os.path.exists(package_sources_path),
+    reason="No aviary.egg-info/SOURCES.txt file, run pip install . to create it",
+)
+def test_egg_sources():
+    with open(package_sources_path) as file:
+        sources = file.read()
+
+    json_files_under_aviary = glob(
+        "**/*.json", recursive=True, root_dir=f"{ROOT}/aviary"
+    )
+
+    for json_file in json_files_under_aviary:
+        assert json_file in sources, f"{json_file} not found in SOURCES.txt"