simstack 0.2.0__py3-none-any.whl → 2.0.2__py3-none-any.whl

simstack/_version.py

@@ -1,5 +1,6 @@
-# file generated by setuptools-scm
+# file generated by vcs-versioning
 # don't change, don't track in version control
+from __future__ import annotations
 
 __all__ = [
     "__version__",
@@ -10,25 +11,14 @@ __all__ = [
     "commit_id",
 ]
 
-TYPE_CHECKING = False
-if TYPE_CHECKING:
-    from typing import Tuple
-    from typing import Union
-
-    VERSION_TUPLE = Tuple[Union[int, str], ...]
-    COMMIT_ID = Union[str, None]
-else:
-    VERSION_TUPLE = object
-    COMMIT_ID = object
-
 version: str
 __version__: str
-__version_tuple__: VERSION_TUPLE
-version_tuple: VERSION_TUPLE
-commit_id: COMMIT_ID
-__commit_id__: COMMIT_ID
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
 
-__version__ = version = '0.2.0'
-__version_tuple__ = version_tuple = (0, 2, 0)
+__version__ = version = '2.0.2'
+__version_tuple__ = version_tuple = (2, 0, 2)
 
 __commit_id__ = commit_id = None

simstack/core/artifacts.py

@@ -1,30 +1,31 @@
-import importlib
 import inspect
 import logging
 import re
-from typing import Optional, List
+from typing import Any, Callable, Optional, List
 
 from odmantic import ObjectId
 
 from simstack.core.context import context
-from simstack.core.engine import current_engine_context
 from simstack.models.artifact_models import ArtifactMapping, ArtifactModel
 from simstack.models.charts_artifact import ChartArtifactModel
 from simstack.models.node_registry import find_child_nodes, NodeRegistry
 from simstack.models.table_artifact import TableArtifactModel
-from simstack.util.importer import function_from_model
-from simstack.util.module_path_checker import is_module_subpath_of_path
+from simstack.util.db import Database
+from simstack.util.importer import import_function
 
 logger = logging.getLogger("artifacts")
 
 
 class ArtifactArguments:
-    def __init__(self, result, task_id: ObjectId = None):
+    def __init__(self, result: Any, task_id: ObjectId | None = None) -> None:
         self.task_id = task_id
         self.result = result
-        self.child_artifacts = []
+        self.child_artifacts: list[ArtifactModel] = []
+        self.call_path: str | None = None
 
-    def add_attributes(self, func, *args, **kwargs):
+    def add_attributes(
+        self, func: Callable[..., Any], *args: Any, **kwargs: Any
+    ) -> None:
         if not callable(func):
             raise ValueError("func must be callable")
         # Get function parameter names
@@ -39,13 +40,12 @@ class ArtifactArguments:
 
 
 async def find_artifact_mappings(
-    node_registry_path: str, task_id: Optional[str] = None
+    node_registry_path: str, db: Database, task_id: Optional[str] = None
 ) -> List[ArtifactMapping]:
     logger.debug(
         f"task_id: {task_id} Loading artifacts with regex pattern: {node_registry_path} "
     )
-    engine = current_engine_context.get()
-    all_mappings = await engine.find(ArtifactMapping)
+    all_mappings = await db.find(ArtifactMapping)
 
     # Filter them manually to find those whose patterns match your path
     matching_mappings = [
@@ -61,9 +61,11 @@ async def find_artifact_mappings(
     return matching_mappings
 
 
-async def register_artifact_mapping(artifact_mapping: ArtifactMapping):
+async def register_artifact_mapping(
+    artifact_mapping: ArtifactMapping,
+) -> ArtifactMapping:
     # Check if an artifact with the same name already exists
-    existing = await context.db.engine.find_one(
+    existing = await context.db.find_one(
         ArtifactMapping, ArtifactMapping.name == artifact_mapping.name
     )
     if existing:
@@ -80,52 +82,12 @@ async def register_artifact_mapping(artifact_mapping: ArtifactMapping):
     ):
         try:
             # Import the function to verify it exists
-            module_path, function_name = artifact_mapping.function_mapping.rsplit(
-                ".", 1
+            func = await import_function(
+                artifact_mapping.function_mapping, context.db, task_id=None
             )
-            module = importlib.import_module(module_path)
-            if module is None:
-                logger.error(f"Module {module_path} could not be imported.")
-                raise ImportError(f"Module {module_path} could not be imported.")
-            # Get the function from the module
-            func = getattr(module, function_name)
-
-            if func:
-                # Function import succeeded, now check if it's in a path with use_pickle=True
-                # Check if the module path is in a non-excluded subdirectory of any path in the path registry
-                for path_name, path_info in context.path_manager.paths.items():
-                    # Check if the path has use_pickle=True
-                    current_module_is_child = is_module_subpath_of_path(
-                        module_path, path_info["path"]
-                    )
-                    if path_info.get("use_pickle", True) and current_module_is_child:
-                        # create a FunctionPickle for the function
-                        try:
-                            from simstack.models.pickle_models import FunctionPickle
-
-                            # Create a FunctionPickle for the function
-                            function_pickle = FunctionPickle(
-                                name=func.__name__, module_path=func.__module__
-                            )
-                            function_pickle.store_function(func)
-
-                            # Save the FunctionPickle to the database
-                            await context.db.save(function_pickle)
-
-                            # Associate the FunctionPickle with the ArtifactMapping
-                            artifact_mapping.pickle_function = function_pickle
-
-                            logger.info(
-                                f"Created FunctionPickle for {artifact_mapping.function_mapping}"
-                            )
-                            break  # We only need to create one FunctionPickle
-                        except Exception as e:
-                            logger.error(
-                                f"Error creating FunctionPickle for {artifact_mapping.function_mapping}: {e}"
-                            )
-            else:
+            if not func:
                 logger.warning(
-                    f"Could not import function {artifact_mapping.function_mapping}"
+                    "Could not import function {artifact_mapping.function_mapping}"
                 )
         except Exception as e:
             logger.error(
@@ -134,26 +96,25 @@ async def register_artifact_mapping(artifact_mapping: ArtifactMapping):
     return await context.db.save(artifact_mapping)
 
 
-async def find_all_artifacts(node_registry: NodeRegistry) -> List[ArtifactModel]:
-    # if not engine:
-    #     engine = context.db.engine
-    engine = current_engine_context.get()
+async def find_all_artifacts(
+    node_registry: NodeRegistry, db: Database
+) -> List[ArtifactModel]:
     return [
-        await engine.find_one(ArtifactModel, ArtifactModel.id == artifact_id)
+        await db.find_one(ArtifactModel, ArtifactModel.id == artifact_id)
         for artifact_id in node_registry.artifact_ids
     ]
 
 
 async def create_artifacts(
     artifact_arguments: ArtifactArguments, node_registry: NodeRegistry
-) -> List[ArtifactModel]:
+) -> List[ObjectId]:
     try:
         call_path = node_registry.call_path
         task_id = node_registry.id
         log_string = f"create artifacts for task_id: {task_id} for {call_path}"
 
         artifact_mappings_list = await find_artifact_mappings(
-            call_path, task_id=task_id
+            call_path, context.db, task_id=task_id
         )
 
         child_nodes = await find_child_nodes(task_id)
@@ -161,9 +122,9 @@ async def create_artifacts(
             f"{log_string} Found {len(artifact_mappings_list)} artifact mappings for path: {call_path} and {len(child_nodes)} child nodes."
         )
         # Concatenate artifacts from all child nodes
-        child_artifacts = []
+        child_artifacts: list[ArtifactModel] = []
         for child_node in child_nodes:
-            loaded_artifacts = await find_all_artifacts(child_node)
+            loaded_artifacts = await find_all_artifacts(child_node, context.db)
             if len(loaded_artifacts) == 1:
                 child_artifacts.extend(loaded_artifacts)
             elif len(loaded_artifacts) > 1:
@@ -182,8 +143,7 @@ async def create_artifacts(
         artifact_arguments.child_artifacts = child_artifacts
         artifact_arguments.call_path = call_path
 
-        artifact_list = []
-        engine = current_engine_context.get()
+        artifact_list: list[ArtifactModel] = []
 
         if len(artifact_mappings_list) > 0:
             for artifact_mapping in artifact_mappings_list:
@@ -196,8 +156,10 @@ async def create_artifacts(
                 )
                 logging.info(log_string_mapping)
                 if artifact_mapping.function_mapping != "CODE":
-                    func = await function_from_model(
-                        artifact_mapping, artifact_arguments.task_id
+                    func = await import_function(
+                        artifact_mapping.function_mapping,
+                        context.db,
+                        artifact_arguments.task_id,
                     )
                     if not func:
                         logger.error(
@@ -237,19 +199,19 @@ async def create_artifacts(
                         continue
                     if isinstance(artifact, TableArtifactModel):
                         artifact.parent_id = node_registry.id
-                        saved_artifact = await engine.save(artifact)
+                        saved_artifact = await context.db.save(artifact)
                         logger.debug(
                             f"{log_string_mapping} new table: {saved_artifact}"
                         )
                     elif isinstance(artifact, ChartArtifactModel):
                         artifact.parent_id = node_registry.id
-                        saved_artifact = await engine.save(artifact)
+                        saved_artifact = await context.db.save(artifact)
                         logger.debug(
                             f"{log_string_mapping} new table: {saved_artifact}"
                         )
                     elif isinstance(artifact, ArtifactModel):
                         artifact.path = call_path
-                        saved_artifact = await engine.save(artifact)
+                        saved_artifact = await context.db.save(artifact)
                         logger.debug(f"{log_string_mapping} new: {saved_artifact}")
                         artifact_list.append(saved_artifact)
                     else:
@@ -284,7 +246,7 @@ async def save_artifact_model(artifact: ArtifactModel) -> ArtifactModel:
         artifact.items = saved_items
 
     # Get the model's collection
-    collection = context.db.engine.get_collection(ArtifactModel)
+    collection = context.db.collection(ArtifactModel)
 
     # Convert model to dict
     item_dict = artifact.model_dump(by_alias=True)
@@ -314,13 +276,13 @@ async def save_artifact_model(artifact: ArtifactModel) -> ArtifactModel:
 
 
 async def find_artifacts(
-    model_class: type[ArtifactModel], **kwargs
+    model_class: type[ArtifactModel], **kwargs: Any
 ) -> List[ArtifactModel]:
     """
     Custom find function for polymorphic models
     """
     # Get the collection
-    collection = context.db.engine.get_collection(ArtifactModel)
+    collection = context.db.collection(ArtifactModel)
 
     # Add item_type filter for subclasses
     if model_class != ArtifactModel:
@@ -378,9 +340,6 @@ async def find_artifacts(
 #     Returns:
 #         An instance of the appropriate ArtifactModel subclass, or None if not found
 #     """
-#     # Get the collection
-#     collection = context.db.engine.get_collection(ArtifactModel)
-#
 #     # Find the document by ID
 #     doc = await collection.find_one({"_id": artifact_id})
 #

simstack/core/config_file.py

@@ -2,7 +2,7 @@ from pathlib import Path
 from typing import Optional
 
 
-class ConfigFile:
+class SimstackConfig:
     """
     Singleton class that manages the path to a configuration file.
 
@@ -10,13 +10,13 @@ class ConfigFile:
     and provides a centralized way to access the configuration file path.
     """
 
-    _instance: Optional["ConfigFile"] = None
+    _instance: Optional["SimstackConfig"] = None
     _config_file_path: Optional[str] = "simstack.toml"  # Default config file name
 
-    def __new__(cls) -> "ConfigFile":
+    def __new__(cls) -> "SimstackConfig":
         """Create or return the singleton instance."""
         if cls._instance is None:
-            cls._instance = super(ConfigFile, cls).__new__(cls)
+            cls._instance = super(SimstackConfig, cls).__new__(cls)
         return cls._instance
 
     def __init__(self):
@@ -50,11 +50,11 @@ class ConfigFile:
 
 
 # Convenience function to get the singleton instance
-def get_config_file() -> ConfigFile:
+def get_config_file() -> SimstackConfig:
     """
     Get the singleton ConfigManager instance.
 
     Returns:
         The ConfigManager singleton instance
     """
-    return ConfigFile()
+    return SimstackConfig()