simstack 2.0.1__py3-none-any.whl

simstack/util/mappings.py

@@ -0,0 +1,41 @@
+from typing import Dict, Optional, List
+from simstack.models.models import ModelMapping, NodeModel
+from odmantic import AIOEngine
+
+class ModelMappingTable:
+    def __init__(self, mappings: List[ModelMapping]):
+        self._by_name: Dict[str, ModelMapping] = {m.name: m for m in mappings}
+        self._by_mapping: Dict[str, ModelMapping] = {m.mapping: m for m in mappings}
+
+    def get_by_name(self, name: str) -> Optional[ModelMapping]:
+        return self._by_name.get(name)
+
+    def get_by_mapping(self, mapping: str) -> Optional[ModelMapping]:
+        return self._by_mapping.get(mapping)
+
+    @classmethod
+    async def load(cls, engine: AIOEngine):
+        mappings = await engine.find(ModelMapping)
+        return cls(list(mappings))
+
+class NodeMappingTable:
+    def __init__(self, nodes: List[NodeModel]):
+        self._by_name: Dict[str, NodeModel] = {n.name: n for n in nodes}
+        self._by_mapping: Dict[str, NodeModel] = {n.function_mapping: n for n in nodes}
+
+    def get_by_name(self, name: str) -> Optional[NodeModel]:
+        return self._by_name.get(name)
+
+    def get_by_mapping(self, mapping: str) -> Optional[NodeModel]:
+        return self._by_mapping.get(mapping)
+
+    @classmethod
+    async def load(cls, engine: AIOEngine):
+        try:
+            nodes = await engine.find(NodeModel)
+        except Exception as e:
+            print(f"Error loading NodeModel from database: {e}")
+            print("WARNING: Returning empty NodeMappingTable to fix obviously corrupted case.")
+            print("Rebuild the NodeMappingTable by rerunning the database initialization or migration scripts.")
+            nodes = []
+        return cls(list(nodes))

simstack/util/minimal_route_finder.py

@@ -0,0 +1,72 @@
+from collections import deque
+from typing import List, Dict
+
+
+def find_minimal_route(
+    routes: List[Dict[str, str]], source: str, target: str
+) -> List[Dict[str, str]]:
+    """
+    Find the minimal (shortest) route from source to target.
+
+    Args:
+        routes: List of route dictionaries with 'source', 'target', and 'host'
+        source: The starting node
+        target: The destination node
+
+    Returns:
+        A list of route dictionaries representing the shortest path from source to target,
+        or an empty list if no path exists
+    """
+    if not routes:
+        return []
+
+    # Build the graph and route map
+    adj: Dict[str, List[str]] = {}
+    route_map: Dict[tuple, Dict[str, str]] = {}
+    nodes = set()
+
+    for route in routes:
+        src = route["source"]
+        dst = route["target"]
+        nodes.add(src)
+        nodes.add(dst)
+
+        if src not in adj:
+            adj[src] = []
+
+        # Avoid adding duplicates to adjacency list, but allow route_map update (last one wins)
+        if dst not in adj[src]:
+            adj[src].append(dst)
+
+        route_map[(src, dst)] = route
+
+    if source not in nodes or target not in nodes:
+        return []
+
+    # BFS to find shortest path
+    queue = deque([(source, [source])])
+    visited = {source}
+
+    while queue:
+        current_node, path = queue.popleft()
+
+        if current_node == target:
+            # Convert the path of nodes to a list of route dictionaries
+            path_routes = []
+            for i in range(len(path) - 1):
+                u = path[i]
+                v = path[i + 1]
+                path_routes.append(route_map[(u, v)])
+            return path_routes
+
+        if current_node in adj:
+            for neighbor in adj[current_node]:
+                if neighbor not in visited:
+                    visited.add(neighbor)
+                    queue.append((neighbor, path + [neighbor]))
+
+    return []
+
+
+# Add an alias for find_minimal_route as find_shortest_route for compatibility
+find_shortest_route = find_minimal_route

simstack/util/mock_db.py

@@ -0,0 +1,144 @@
+import asyncio
+from typing import Type, TypeVar, Dict, List, Optional
+
+from odmantic import Model, ObjectId
+
+T = TypeVar("T", bound=Model)
+
+
+class InMemoryCollection:
+    def __init__(self, engine: "InMemoryAIOEngine"):
+        self._engine = engine
+        self._store: List[Dict] = []
+
+    def __getitem__(self, index: int) -> Dict:
+        return self._store[index]
+
+    def __setitem__(self, index: int, value: Dict):
+        self._store[index] = value
+
+    def __iter__(self):
+        return iter(self._store)
+
+
+class InMemoryAIOEngine:
+    def __init__(self):
+        # Stores data as {collection_name: [documents]}
+        self._store: Dict[str, InMemoryCollection] = {}
+
+    async def save(self, instance: T) -> T:
+        collection_name = (
+            instance.__collection__
+            if hasattr(instance, "__collection__")
+            else instance.__class__.__name__.lower()
+        )
+
+        obj_dict = instance.model_dump()
+
+        # Ensure the instance has an ObjectId assigned
+        if not instance.id:
+            instance.id = ObjectId()
+            obj_dict["id"] = instance.id
+
+        # Initialize collection if it doesn't exist
+        if collection_name not in self._store:
+            self._store[collection_name] = InMemoryCollection(self)
+
+        # Update if already exists
+        existing_index = next(
+            (
+                i
+                for i, obj in enumerate(self._store[collection_name])
+                if obj["id"] == instance.id
+            ),
+            None,
+        )
+        if existing_index is not None:
+            self._store[collection_name][existing_index] = obj_dict
+        else:
+            self._store[collection_name].append(obj_dict)
+
+        return instance
+
+    async def find(self, model: Type[T], **query) -> List[T]:
+        collection_name = (
+            model.__collection__
+            if hasattr(model, "__collection__")
+            else model.__name__.lower()
+        )
+        all_objects = self._store.get(collection_name, [])
+        results = []
+
+        for obj_dict in all_objects:
+            if all(obj_dict.get(k) == v for k, v in query.items()):
+                results.append(model.model_validate(obj_dict))
+        return results
+
+
+async def find_one(self, model: Type[T], condition=None, **query) -> Optional[T]:
+    # Handle both keyword arguments and condition objects
+    if condition is not None:
+        # For conditions like ModelMapping.name == class_name
+        # This is a more robust approach to extract field name and value
+        condition_str = str(condition)
+
+        # Look for the pattern: field_name == value or field_name: {$eq: value}
+        if "==" in condition_str:
+            # Handle direct comparison like ModelMapping.name == 'TestClass'
+            parts = condition_str.split("==")
+            if len(parts) == 2:
+                field_name = parts[0].strip().split(".")[-1]
+                value = parts[1].strip().strip("'\"")
+                query[field_name] = value
+        elif "$eq" in condition_str:
+            # Handle MongoDB-style query like {'mapping': {'$eq': 'tests.core.test_import_class.TestClass'}}
+            import re
+
+            # Extract field name and value using regex
+            field_match = re.search(
+                r"'(\w+)':\s*\{'?\$eq'?:\s*'([^']+)'", condition_str
+            )
+            if field_match:
+                field_name = field_match.group(1)
+                value = field_match.group(2)
+                query[field_name] = value
+            else:
+                # Fallback - try to extract from the string representation
+                # Look for pattern like "ModelMapping.field_name"
+                field_match = re.search(r"(\w+)\.(\w+)", condition_str)
+                if field_match:
+                    field_name = field_match.group(2)
+                    # Try to extract the value (this is still a fallback)
+                    value_match = re.search(r"'([^']+)'(?:[^']*$)", condition_str)
+                    if value_match:
+                        value = value_match.group(1)
+                        query[field_name] = value
+
+    found = await self.find(model, **query)
+    return found[0] if found else None
+
+
+# Example Usage:
+if __name__ == "__main__":
+    from odmantic import Model
+
+    class BinaryOperationInput(Model):
+        arg1: int
+        arg2: int
+
+    async def test_engine():
+        engine = InMemoryAIOEngine()
+        obj = BinaryOperationInput(arg1=1, arg2=2)
+
+        # Saving object
+        await engine.save(obj)
+
+        # Finding objects
+        found_objs = await engine.find(BinaryOperationInput, arg1=1)
+        print("Found objects:", found_objs)
+
+        # Find one object
+        found_one = await engine.find_one(BinaryOperationInput, arg2=2)
+        print("Found single object:", found_one)
+
+    asyncio.run(test_engine())

simstack/util/module_path_checker.py

@@ -0,0 +1,38 @@
+import os
+from pathlib import Path
+
+from simstack.util.project_root_finder import find_project_root
+
+
+def is_module_subpath_of_path(module_path: str, path_info_path: Path) -> bool:
+    """
+    Check whether a dot-separated module path is a subpath of an absolute path.
+
+    Args:
+        module_path: A dot-separated path relative to the project root (e.g., "src.simstack.core.artifacts")
+        path_info_path: An absolute path below the project root (e.g., "/path/to/project/src/simstack/core")
+
+    Returns:
+        True if the module path corresponds to a location under the given absolute path, False otherwise
+    """
+
+    # Convert the module path to the file system path
+    # Replace dots with path separators
+    module_file_path = module_path.replace(".", os.sep)
+
+    # Convert path_info_path to the Path object for easier manipulation
+    path_info_path_obj = Path(path_info_path)
+
+    # Get the project root using the existing utility function
+    project_root = find_project_root()
+
+    # Convert the module path to the absolute path relative to the project root
+    module_absolute_path = project_root / module_file_path
+
+    try:
+        # Check if the module path is under the given path_info_path
+        module_absolute_path.relative_to(path_info_path_obj)
+        return True
+    except ValueError:
+        # If relative_to raises ValueError, the module path is not under path_info_path
+        return False

simstack/util/object_list_mixin.py

@@ -0,0 +1,282 @@
+import re
+from typing import Generic, Optional, List, Iterator, Type, Union, Iterable, TypeVar, TYPE_CHECKING
+
+from odmantic import ObjectId, Model
+
+from simstack.util.generic_list_mixin import GenericListMixin
+
+if TYPE_CHECKING:
+    from simstack.util.db import Database
+
+T = TypeVar("T")
+
+
+class ObjectListMixin(GenericListMixin[ObjectId], Generic[T]):
+    """
+       Mixin class for lists of Model ObjectIDs.
+       Stores ObjectId instances in `elements`, but allows interaction with Model instances.
+       """
+
+    @staticmethod
+    def _normalize_elements_for_init(data: dict) -> tuple[dict, Optional[List[Model]]]:
+        elements = data.get("elements")
+        if not elements:
+            return data, None
+
+        elements = list(elements)
+        contains_models = any(isinstance(element, Model) for element in elements)
+
+        if not contains_models:
+            return data, None
+
+        if not all(isinstance(element, Model) for element in elements):
+            raise ValueError("elements must contain either only model instances or only ObjectIds")
+
+        normalized_data = dict(data)
+        normalized_data["elements"] = [element.id for element in elements]
+        return normalized_data, elements
+
+    def __iter__(self) -> Iterator[T]:
+        return iter(self._get_cache())
+
+    def _get_model_class(self) -> Type[T]:
+        # T is the first type argument of ObjectListMixin
+        if hasattr(self, "__orig_bases__"):
+            for base in self.__class__.__orig_bases__:
+                if hasattr(base, "__origin__") and base.__origin__ is ObjectListMixin:
+                    return base.__args__[0]
+        raise RuntimeError(f"Could not determine model class for {self.__class__.__name__}")
+
+    def _set_cache(self, cache: List[T]):
+        object.__getattribute__(self, "__dict__")["_cache"] = cache
+        return cache
+
+    def _get_cache(self) -> List[T]:
+        self_dict = object.__getattribute__(self, "__dict__")
+        cache = self_dict.get("_cache",None)
+        if cache is None:
+            cache = []
+            cache = self._set_cache(cache)
+        return cache
+
+
+    def append(self, element: T):
+        if not isinstance(element, Model):
+            raise ValueError("can only append models to ObjectListMixin")
+        cache = self._get_cache()
+        if element not in cache:
+            cache.append(element)
+        self._set_cache(cache)
+
+        obj_id = getattr(element, "id", None)
+        if obj_id is None:
+            raise ValueError("attempting to append a model without id to an ObjectListMixin class")
+        if obj_id not in self.elements:
+            self.elements.append(obj_id)
+
+    def extend(self, elements: Union[List[T], "ObjectListMixin[T]"]):
+        for element in elements:
+           self.append(element)
+
+    async def save(self, db: "Database"):
+        cache = self._get_cache()
+        # Save each element first (this might trigger their own custom saves)
+        for element in cache:
+            await db.save(element)
+        if isinstance(self, Model):  # embedded models are save with the parent
+            await db.save_unchecked(self)
+        return self
+
+    def insert(self, index: int, element: T):
+        self.elements.insert(index, element.id)
+        cache = self._get_cache()
+        if index <= len(cache):
+            cache.insert(index, element)
+
+    def remove(self, element: T):
+        self.elements.remove(element.id)
+        cache = self._get_cache()
+        if element in cache:
+            cache.remove(element)
+
+    def pop(self, index: int = -1) -> T:
+        self.elements.pop(index)
+        cache = self._get_cache()
+        if index < 0:
+            index = len(self.elements) + 1 + index # elements already popped
+        if index < len(cache):
+            return cache.pop(index)
+        else:
+            # If it's not in cache, we might have a problem if it was never loaded.
+            # But pop is usually used on loaded lists.
+            raise IndexError("Index out of range for cache. Load elements first.")
+
+    def get(self, index: int) -> T:
+        """Get an element by index, loading it from DB if necessary."""
+        if index < 0 or index >= len(self.elements):
+            raise IndexError("list index out of range")
+
+        cache = self._get_cache()
+        if index < len(cache):
+            return cache[index]
+
+        raise IndexError("Index out of range for cache. Load elements first.")
+
+    def count(self, element: T) -> int:
+        return self.elements.count(element.id)
+
+    def index(self, element: T, start: int = 0, stop: int = None) -> int:
+        return self.elements.index(element.id, start, stop if stop is not None else len(self.elements))
+
+    def reverse(self):
+        self.elements.reverse()
+        self._get_cache().reverse()
+
+    def sort(self, key=None, reverse: bool = False):
+        # This is tricky because we need to sort both.
+        # Simplest is to sort the cache and then update elements.
+        # But we might not have a full cache.
+        if len(self._get_cache()) != len(self.elements):
+            # We can't easily sort if we don't have all elements in cache.
+            # For now, let's just sort the elements if key only depends on ObjectId,
+            # or raise if it needs the Model.
+            # Actually, standard list.sort() uses the objects.
+            raise RuntimeError("Sort requires all elements to be loaded in cache. Call _load_all_to_cache() first.")
+
+        cache = self._get_cache()
+        cache.sort(key=key, reverse=reverse)
+        self.elements[:] = [obj.id for obj in cache]
+
+    def __setitem__(self, index: Union[int, slice], value: Union[T, List[T]]):
+        if isinstance(index, slice):
+            if isinstance(value, Iterable):
+                ids = [v.id for v in value]
+                self.elements[index] = ids
+                cache = self._get_cache()
+                # Slice assignment on cache is only safe if it matches the current loaded state
+                # For simplicity, let's just clear cache or try to update if it fits
+                if index.start is not None and index.stop is not None and index.stop <= len(cache):
+                     cache[index] = list(value)
+                self._set_cache(cache)
+            else:
+                raise TypeError("Can only assign an iterable to a slice")
+        else:
+            self.elements[index] = value.id
+            cache = self._get_cache()
+            if index < len(cache):
+                cache[index] = value
+            self._set_cache(cache)
+
+    def __delitem__(self, index: Union[int, slice]):
+        del self.elements[index]
+        cache = self._get_cache()
+        if isinstance(index, slice):
+            if index.start is not None and index.stop is not None and index.stop <= len(cache):
+                del cache[index]
+            else:
+                raise IndexError("Slice out of range for cache")
+        else:
+            if index < len(cache):
+                del cache[index]
+        self._set_cache(cache)
+
+    def copy(self) -> "ObjectListMixin[T]":
+        import copy
+        return copy.copy(self)
+
+    def __getitem__(self, index: Union[int, slice]) -> Union[T, List[T], ObjectId, List[ObjectId]]:
+        cache = self._get_cache()
+        if isinstance(index, slice):
+            if index.start is not None and index.stop is not None and index.stop <= len(cache):
+                return cache[index]
+            return self.elements[index]
+        else:
+            if 0 <= index < len(cache):
+                return cache[index]
+            return self.elements[index]
+
+    def find(self, pattern: str) -> Optional[T]:
+        regex = re.compile(pattern)
+        cache = self._get_cache()
+        for obj in cache:
+            name = getattr(obj, "name", None)
+            if name and regex.search(name):
+                return obj
+        return None
+
+    def find_all(self, pattern: str) -> Iterator[T]:
+        regex = re.compile(pattern)
+        cache = self._get_cache()
+        for obj in cache:
+            name = getattr(obj, "name", None)
+            if name and regex.search(name):
+                yield obj
+
+
+    async def db_find_postprocess(self, db: "Database"):
+        """Instance-level post-processing"""
+        cache = self._get_cache()
+        model_class = self._get_model_class()
+        # Identify which IDs are missing from the cache
+        missing_ids = [obj_id for obj_id in self.elements if not any(getattr(o, "id", None) == obj_id for o in cache)]
+
+        if missing_ids:
+            # Load all missing elements at once
+            query = model_class.id.in_(missing_ids)
+            # engine.find might return a cursor or a list
+            results = await db.find(model_class, query)
+
+            # Map results by ID for efficient lookup
+            loaded_map = {}
+            for obj in results:
+                # Recursively unwrap nested collections (mock DB behavior)
+                it = obj
+                while isinstance(it, (list, tuple)) and len(it) > 0:
+                    it = it[0]
+
+                # Check for Model or duck-typed object
+                if hasattr(it, "id"):
+                    loaded_map[it.id] = it
+
+            # Update cache maintaining the order of self.elements
+            new_cache = []
+            for obj_id in self.elements:
+                # Find in existing cache
+                existing = next((o for o in cache if getattr(o, "id", None) == obj_id), None)
+                if existing:
+                    new_cache.append(existing)
+                elif obj_id in loaded_map:
+                    new_cache.append(loaded_map[obj_id])
+                else:
+                    import logging
+                    logging.getLogger(__name__).warning(f"Could not load object {obj_id} for {model_class.__name__}")
+            self._set_cache(new_cache)
+        return self
+
+    def delete_element(self, element: Union[T, ObjectId]):
+        cache = self._get_cache()
+        if isinstance(element, Model):
+            obj_id = element.id
+            if element in cache:
+                cache.remove(element)
+        else:
+            obj_id = element
+            # If it's an ObjectId, we might need to find the object in cache to remove it
+            cache = [o for o in cache if getattr(o, "id", None) != obj_id]
+        self._set_cache(cache)
+        if obj_id in self.elements:
+            self.elements.remove(obj_id)
+
+    def clear(self):
+        self.elements.clear()
+        self._set_cache([])
+
+    def __contains__(self, element: Union[T, ObjectId]) -> bool:
+        if isinstance(element, Model):
+            return element.id in self.elements
+        elif isinstance(element, ObjectId):
+            return element in self.elements
+        return element in self._get_cache()
+
+    def __len__(self) -> int:
+        return len(self.elements)