Merge pull request #92 from SWIFTSIM/multiple-registration-files

Added support for multiple registration files.

Author: bwvdnbro

velociraptor/__init__.py

@@ -30,13 +30,13 @@
 from velociraptor.catalogue.catalogue import VelociraptorCatalogue
 from velociraptor.__version__ import __version__
 
-from typing import Union
+from typing import Union, List
 
 
 def load(
     filename: str,
     disregard_units: bool = False,
-    registration_file_path: Union[str, None] = None,
+    registration_file_path: Union[List[str], str, None] = None,
     mask: slice = Ellipsis,
 ) -> VelociraptorCatalogue:
     """
@@ -59,8 +59,8 @@ def load(
         star formation rate units are presented in non-internal
         units.
 
-    registration_file_path: str, optional
-        The filename of the derived quantities script to register
+    registration_file_path: Union[List[str], str], optional
+        The filename of the derived quantities script(s) to register
         additional properties with the catalogue. This is an
         advanced feature. See the documentation for more details.
 

velociraptor/catalogue/catalogue.py

@@ -410,7 +410,9 @@ def __create_sub_catalogues(self):
 
         return
 
-    def register_derived_quantities(self, registration_file_path: str) -> None:
+    def register_derived_quantities(
+        self, registration_file_path: Union[List[str], str]
+    ) -> None:
         """
         Register any required derived quantities. These will
         be available through `catalogue.derived_quantities.{your_names}`.
@@ -421,8 +423,8 @@ def register_derived_quantities(self, registration_file_path: str) -> None:
         Parameters
         ----------
 
-        registration_file_path: str
-            Path to the python source file that contains the code to
+        registration_file_path: Union[List[str], str]
+            Path to the python source file(s) that contain(s) the code to
             register the additional derived quantities.
         """
 

velociraptor/catalogue/derived.py

@@ -5,20 +5,21 @@
 
 import numpy as np
 import unyt
+from typing import Union, List
 
 
 class DerivedQuantities(object):
     """
-    Derived quantities class. Contains methods to open a python
-    source file that creates derived quantities as follows:
+    Derived quantities class. Contains methods to open (a) python
+    source file(s) that create(s) derived quantities as follows:
 
     The source file will have access to:
 
     + numpy (imported as np)
     + unyt (imported as unyt)
 
     You should write your derived quantities as follows:
-    
+
     .. code-block:: python
 
         self.derived_quantitiy_name = catalogue.type.field_name * 0.5
@@ -49,9 +50,9 @@ class DerivedQuantities(object):
             setattr(
                 self, f"specific_sfr_gas_{aperture_size}_kpc", ssfr
             )
-    
 
-    The path to this file should be passed to the __init__ method
+
+    The path to this file(s) should be passed to the __init__ method
     of this class. Note that you should only register quantities that
     you do in fact intend to use, as these are not lazily loaded
     in the same way as the properties that are built into catalogues.
@@ -61,36 +62,39 @@ class DerivedQuantities(object):
 
     """
 
-    def __init__(self, registration_file_path: str, catalogue):
+    def __init__(self, registration_file_path: Union[List[str], str], catalogue):
         """
-        Registers additional (derived) quantities from the 
+        Registers additional (derived) quantities from the
         VelociraptorCatalogue to itself, using a python
         source file that does this registration inside
         the private _register_quantities() method.
 
         Parameters
         ----------
 
-        registration_file_path: str
-            Path to the python source file. For more information
-            on the contents of this file, check out the information
+        registration_file_path: Union[List[str], str]
+            Path to the python source file(s). For more information
+            on the contents of this file/these files, check out the information
             of this object.
 
         catalogue: VelociraptorCatalogue
             The catalogue to derive the quantities from.
 
-        
+
         Returns
         -------
-    
+
         DerivedQuantities
             An instance of the DerivedQuantities class with
             the properties defined in registration_file_path
             available as attributes.
-        
+
         """
 
-        self.registration_file_path = registration_file_path
+        if isinstance(registration_file_path, list):
+            self.registration_file_paths = list(registration_file_path)
+        else:
+            self.registration_file_paths = [registration_file_path]
         self.catalogue = catalogue
 
         self._register_quantities()
@@ -105,8 +109,8 @@ def _register_quantities(self):
         """
 
         catalogue = self.catalogue
-
-        with open(self.registration_file_path, "r") as handle:
-            exec(handle.read())
+        for file_path in self.registration_file_paths:
+            with open(file_path, "r") as handle:
+                exec(handle.read())
 
         return