update pytest config

Author: bmtgoncalves

docs/conf.py

@@ -1,6 +1,7 @@
 import os
 import sys
-sys.path.insert(0, os.path.abspath('../src/epidemik/'))
+
+sys.path.insert(0, os.path.abspath("../src/epidemik/"))
 import epidemik
 
 # Configuration file for the Sphinx documentation builder.
@@ -11,45 +12,44 @@
 # -- Project information -----------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
 
-project = 'epidemik'
-copyright = '2024, Bruno Gonçalves'
-author = 'Bruno Gonçalves'
+project = "epidemik"
+copyright = "2024, Bruno Gonçalves"
+author = "Bruno Gonçalves"
 release = epidemik.__version__
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
 
 extensions = [
-    'sphinx_rtd_theme',
-    'sphinx.ext.duration',
-    'sphinx.ext.doctest',
-    'sphinx.ext.autodoc',
-    'sphinx.ext.autosummary',
+    "sphinx_rtd_theme",
+    "sphinx.ext.duration",
+    "sphinx.ext.doctest",
+    "sphinx.ext.autodoc",
+    "sphinx.ext.autosummary",
 ]
 
-templates_path = ['_templates']
-exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']
-
+templates_path = ["_templates"]
+exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
 
 
 # -- Options for HTML output -------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
 
-html_theme = 'sphinx_rtd_theme'
-html_static_path = ['_static']
+html_theme = "sphinx_rtd_theme"
+html_static_path = ["_static"]
 html_theme_options = {
-    'analytics_id': 'G-HKWS10TRJ1', 
-    'analytics_anonymize_ip': False,
-    'logo_only': False,
-    'display_version': True,
-    'prev_next_buttons_location': 'bottom',
-    'style_external_links': True,
-    'vcs_pageview_mode': '',
-    'style_nav_header_background': 'white',
+    "analytics_id": "G-HKWS10TRJ1",
+    "analytics_anonymize_ip": False,
+    "logo_only": False,
+    "display_version": True,
+    "prev_next_buttons_location": "bottom",
+    "style_external_links": True,
+    "vcs_pageview_mode": "",
+    "style_nav_header_background": "white",
     # Toc options
-    'collapse_navigation': True,
-    'sticky_navigation': True,
-    'navigation_depth': 4,
-    'includehidden': True,
-    'titles_only': False
-}
+    "collapse_navigation": True,
+    "sticky_navigation": True,
+    "navigation_depth": 4,
+    "includehidden": True,
+    "titles_only": False,
+}

poetry.lock

@@ -33,6 +33,7 @@ enable = true
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.3.4"
+pytest-cov = "^6.0.0"
 
 [build-system]
 requires = ["poetry-core>=2.0.0", "poetry-dynamic-versioning>=1.0.0,<2.0.0"]

pyproject.toml

@@ -0,0 +1,2 @@
+[pytest]
+addopts = --ignore=.venv --cov=epidemik

pytest.ini

@@ -3,7 +3,7 @@
 # @author Bruno Goncalves
 ######################################################
 
-from typing import Dict, List, Set, Union
+from typing import Dict, List, Set, Union, Self
 import warnings
 import string
 import time
@@ -789,6 +789,33 @@ def save_model(self, filename: str) -> None:
         with open(filename, "wt") as f:
             f.write(self.__repr__())
 
+    def list_models() -> List[str]:
+        """
+        List the models available in the official repository
+        """
+        remote_path = utils.get_remote_path()
+        remote_path = os.path.join(remote_path, "model_list.txt")
+
+        cache_dir = utils.get_cache_directory()
+
+        if not os.path.exists(cache_dir):
+            os.makedirs(cache_dir)
+
+        local_path = os.path.join(cache_dir, "model_list.txt")
+
+        # Always get the latest version
+        urlretrieve(remote_path, local_path)
+        models = [line.strip() for line in open(local_path, "rt").readlines()]
+
+        # Add any user models
+        if os.path.exists(cache_dir):
+            for file in os.listdir(cache_dir):
+                if file.endswith(".yaml"):
+                    models.append(file)
+
+        # Make sure to remove duplicates
+        return sorted(set(models))
+
     def load_model(filename: str) -> None:
         """
         Load the model from a file
@@ -824,26 +851,24 @@ def load_model(filename: str) -> None:
 
         return model
 
-
-    def download_model(filename: str, repo: Union[str, None] = None) -> None:
+    def download_model(
+        filename: str, repo: Union[str, None] = None, load_model: bool = True
+    ) -> Union[None, Self]:
         """
         Download model from offical repository
         """
-        if repo is None:
-            repo = utils.OFFICIAL_REPO
-
-        parsed_repo = urlparse(repo)
+        remote_path = utils.get_remote_path(repo)
+        remote_path = os.path.join(remote_path, filename)
 
-        if parsed_repo.netloc == 'github.com':
-            repo = repo.replace('github.com', 'raw.githubusercontent.com')
-            remote_path = repo + os.path.join('refs/heads/models/models/', filename)
+        cache_dir = utils.get_cache_directory()
 
-        if not os.path.exists(utils.LOCAL_DIRECTORY):
-            os.makedirs(utils.LOCAL_DIRECTORY)
+        if not os.path.exists(cache_dir):
+            os.makedirs(cache_dir)
 
-        local_path = os.path.join(utils.LOCAL_DIRECTORY, filename)
+        local_path = os.path.join(cache_dir, filename)
 
         if not os.path.exists(local_path):
             urlretrieve(remote_path, local_path)
 
-        return EpiModel.load_model(local_path)
+        if load_model:
+            return EpiModel.load_model(local_path)

src/epidemik/EpiModel.py

@@ -5,6 +5,8 @@
 
 import networkx as nx
 import numpy as np
+import time
+import os
 from numpy import linalg
 from numpy import random
 import pandas as pd
@@ -14,7 +16,8 @@
 
 from typing import Union
 
-from .EpiModel import *
+from .EpiModel import EpiModel
+from . import utils
 
 from tqdm import tqdm
 
@@ -214,7 +217,7 @@ def simulate(
         **kwargs,
     ) -> None:
         if seed_state is None:
-            raise NotInitialized("You have to specify the seed_state")
+            raise utils.NotInitialized("You have to specify the seed_state")
 
         self._initialize_populations(susceptible)
 
@@ -233,7 +236,7 @@ def simulate(
             self.compartments_ = self._run_travel(self.compartments_, self.travel_graph)
 
     def integrate(self, **kwargs):
-        raise NotImplementedError(
+        raise utils.NotImplementedError(
             "MetaEpiModel doesn't support direct integration of the ODE"
         )