import mergertrees and coretrees modules in main __init__.py, fix pat…

…h-simulation argument in readers
ArgonneCPAC · Jun 11, 2024 · 1b9af0c · 1b9af0c
1 parent 2b06a84
commit 1b9af0c
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 2 deletions.
diff --git a/haccytrees/__init__.py b/haccytrees/__init__.py
@@ -5,5 +5,7 @@
 __version__ = "0.9.1"
 
 from .simulations import Simulation
+from . import mergertrees
+from . import coretrees
 
 __all__ = ["Simulation"]
diff --git a/haccytrees/coretrees/coretree_reader.py b/haccytrees/coretrees/coretree_reader.py
@@ -153,7 +153,7 @@ def corematrix_reader(
     calculate_secondary_host_row: bool = False,
 ):
     if isinstance(simulation, str):
-        if simulation[:-4] == ".cfg":
+        if simulation[-4:] == ".cfg":
             simulation = Simulation.parse_config(simulation)
         else:
             simulation = Simulation.simulations[simulation]

diff --git a/haccytrees/mergertrees/forest_reader.py b/haccytrees/mergertrees/forest_reader.py
@@ -208,7 +208,7 @@ def read_forest(
 
     """
     if isinstance(simulation, str):
-        if simulation[:-4] == ".cfg":
+        if simulation[-4:] == ".cfg":
             simulation = Simulation.parse_config(simulation)
         else:
             simulation = Simulation.simulations[simulation]