simstack 0.2.0__py3-none-any.whl

simstack/core/context.py

@@ -0,0 +1,167 @@
+import logging  # Import logging before using it
+import os
+import sys
+from urllib.parse import urlparse, urlunparse
+
+from simstack.core.definitions import DBType
+from simstack.util.config_reader import ConfigReader
+from simstack.util.path_manager import PathManager
+
+# from simstack.core.model_table import make_models_for_path
+# from simstack.core.node_table import make_nodes_for_path
+from simstack.util.project_root_finder import find_project_root
+from simstack.util.setup_logging import setup_logging
+
+
+def remove_password_from_connection_string(connection_string):
+    parsed_url = urlparse(connection_string)
+
+    # Extract username and rebuild netloc without password
+    netloc = parsed_url.hostname
+    if parsed_url.username:
+        netloc = f"{parsed_url.username}@{netloc}"
+    if parsed_url.port:
+        netloc += f":{parsed_url.port}"
+
+    clean_url = parsed_url._replace(netloc=netloc)
+
+    return urlunparse(clean_url)
+
+
+class GlobalState:
+    _instance = None
+    _initialized = False
+
+    def __new__(cls, *args, **kwargs):
+        if not cls._instance:
+            cls._instance = super().__new__(cls)
+            # TODO the DB used in simstack applications is fixed and comes from the config file, the db in the server changes by user
+            cls._instance.db = None
+            cls._instance.log_handler = None
+            cls._instance.path_manager = None
+            cls._instance.config = None
+
+        return cls._instance
+
+    def __init__(self, **kwargs):
+        """Initialize GlobalState instance.
+
+        This method only runs once due to the singleton pattern.
+        Use the initialize() method to set up the instance with database settings.
+        """
+        if not hasattr(self, "_initialized"):
+            self._initialized = True
+
+            self.db = None
+            self.log_handler = None
+            self.path_manager = None
+            self.config = None
+
+            self.initialize(**kwargs)
+
+    def __getattribute__(self, name):
+        # These special attributes should always be accessible
+        if name in ("_initialized", "initialize", "initialized"):
+            return object.__getattribute__(self, name)
+
+        # For other attributes, check initialization
+        if not object.__getattribute__(self, "_initialized"):
+            raise RuntimeError(
+                "GlobalState must be initialized with database settings before use"
+            )
+
+        return object.__getattribute__(self, name)
+
+    # @async_helper
+    # async def remake_models_and_nodes(self,path: str):
+    #     if path is not None:  # rescan all files in the path if pickling is needed
+    #         parent_path = self.path_manager.find_parent_path(path)
+    #         if not parent_path:
+    #             raise ValueError(f"Path '{path}' not found in paths. Please check your configuration.")
+    #         await make_models_for_path(parent_path, self.path_manager, context.db.engine)
+    #         await make_nodes_for_path(parent_path, self.path_manager, context.db.engine)
+
+    def initialize(self, **kwargs):
+        db_name = kwargs.get("db_name", None)
+        connection_string = kwargs.get("connection_string", None)
+        resource_str: str = kwargs.get("resource", "self")
+
+        """Initialize the GlobalState with database settings"""
+        if self._initialized:
+            raise RuntimeError("GlobalState already initialized")
+        self._initialized = True
+
+        is_test = kwargs.get("is_test", False)
+        self.config = ConfigReader(
+            db_name=db_name,
+            connection_string=connection_string,
+            resource=resource_str,
+            is_test=is_test,
+        )
+
+        # check that the database can be reached and set logging up
+        from simstack.util.db import Database
+
+        # Use in-memory database for tests
+        db_type = DBType.IN_MEMORY if is_test and db_name is None else DBType.MONGODB
+
+        try:
+            self.db = Database(
+                db_type, self.config.database_name, self.config.connection_string
+            )
+            if db_type == DBType.MONGODB:
+                # Only ping real MongoDB connections
+                self.db.client.admin.command("ping")
+        except ConnectionError as e:
+            if not is_test:
+                print(f"Could not connect to the database: {e}")
+                sys.exit(-1)
+            else:
+                # For tests, continue without the database connection failure
+                print(f"Warning: Database connection failed in test mode: {e}")
+
+        if is_test:
+            # For tests, use simple console logging without the database handler
+            logging.basicConfig(
+                level=logging.ERROR,
+                format="%(asctime)s - %(name)-15s - %(levelname)-10s - %(filename)-20s:%(lineno)4d - %(message)s",
+            )
+            self.log_handler = logging.getLogger()
+        else:
+            self.log_handler = setup_logging(
+                self.config.connection_string,
+                self.config.database_name,
+                kwargs.get("log_level", "INFO"),
+            )
+
+        # initialize the rest of the variables in the config, but now we can get the errors in the database
+        self.config.secondary_init(kwargs.get("workdir", None))
+
+        logger = logging.getLogger("Context")
+        safe_connection_string = remove_password_from_connection_string(
+            self.config.connection_string
+        )
+        logger.info(
+            f"Database connection established to {self.config.database_name} at {safe_connection_string}"
+        )
+
+        # Initialize PathManager from config
+        self.path_manager = PathManager.from_config(self.config)
+
+        # Set any additional parameters
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+        return self
+
+    @property
+    def initialized(self):
+        return self._initialized
+
+
+root_dir = find_project_root()
+path_dir = os.path.join(root_dir, "src")
+if path_dir not in os.sys.path:
+    os.sys.path.append(path_dir)
+
+# Create the singleton instance, but it's not initialized yet
+context = GlobalState()

simstack/core/data.py

@@ -0,0 +1,139 @@
+import inspect
+import uuid
+from typing import get_origin, get_args, Any, Type, Optional
+
+from pydantic import BaseModel
+from sqlmodel import SQLModel, Field
+
+
+def get_recursive(data: Any, key: str):
+    elements = key.split(".")
+    if len(elements) == 1:
+        if key == "$$":
+            return data
+        else:
+            return data.__dict__[key]
+    else:
+        return get_recursive(data.__dict__[elements[0]], ".".join(elements[1:]))
+
+
+class Assignment(BaseModel):
+    """
+    Represents a mapping container class inherited from BaseModel.
+
+    This class stores a default dictionary mapping and provides a callable
+    method that enables operations using a source data and a target class.
+    It will be used to generate instances of workflow nodes
+    or workflow results from the workflow data
+
+    In the mapping dict the keys specify the target attribute names of the target class.
+
+    If the target class is None, the special key "$$" is used to specify that value.
+    The value specifies the source attribute name of the source data. The special value "$$" means
+    that the value is the source data itself will be used.
+
+    This mechanism allows to map non-class types to each other and to create instances of classes
+    where needed.
+
+    For BaseModel classes the assignment is driven by the model_fields of the target class. For fields that
+    are BaseModel classes, the assignment is recursive.
+
+    :ivar mapping: Stores the dictionary for mapping values between the
+        source and the target class.
+    :type mapping: dict
+    """
+
+    # TODO doe SQLModel for dict
+    id: int = Field(default=None, primary_key=True)
+    # mapping: dict = Field(default_factory=dict)
+
+    model_config = {
+        "frozen": True,
+    }
+
+    def __init__(self, mapping: dict):
+        super().__init__(mapping=mapping)
+
+    def __call__(self, target_class: Type[SQLModel] | None, source: Any):
+        """
+        Create a new instance of the class from the source data
+
+        :param source:
+        :return:
+        """
+        if target_class is None:
+            return get_recursive(source, self.mapping["$$"])
+        else:
+            return target_class.from_data(source, self)
+
+
+def is_data_type(annotation):
+    """Check if a type annotation is or contains Data"""
+    # Direct subclass check
+    if inspect.isclass(annotation) and issubclass(annotation, Data):
+        return True
+
+    # Handle Union types (like Optional[Type])
+    origin = get_origin(annotation)
+    if origin is not None:
+        # For union types (X | Y or Union[X, Y])
+        args = get_args(annotation)
+        return any(is_data_type(arg) for arg in args)
+
+    return False
+
+
+class Data(SQLModel, table=False):
+    id: Optional[uuid.UUID] = Field(default=uuid.uuid4(), primary_key=True)
+    # item_type: str = Field(sa_column_kwargs={"nullable": False})
+    #
+    # __mapper_args__ = {
+    #     "polymorphic_on": "item_type",
+    #     "polymorphic_identity": "base_item",
+    # }
+
+    # TODO automatically set the type to the name of the class
+    # @validator(mode='before')
+    # def set_type_to_class_name(cls, data):
+    #     if isinstance(data, dict):
+    #         # Don't override if explicitly provided
+    #         if 'type' not in data:
+    #             data['type'] = cls.__name__
+    #     return data
+
+    @classmethod
+    def from_data(cls, data: SQLModel, assignment: Assignment):
+        value_dict = {}
+        mapping = assignment.mapping
+        for key, field_info in cls.model_fields.items():
+            if key == "type":
+                value_dict["type"] = cls.__name__
+            else:
+                annotation = field_info.annotation
+                # Check if this field is a Data subclass
+                if is_data_type(annotation) and key in mapping:
+                    # If mapping is a dict, it's a nested mapping for a Data subclass
+                    if isinstance(mapping[key], dict):
+                        # Get the concrete subclass to use
+                        if inspect.isclass(annotation) and issubclass(annotation, Data):
+                            nested_class = annotation
+                        else:
+                            # TODO Handle Union types - this is simplified, might need enhancement
+                            for arg in get_args(annotation):
+                                if inspect.isclass(arg) and issubclass(arg, Data):
+                                    nested_class = arg
+                                    break
+
+                        # Create an instance using the nested mapping
+                        # TODO handle error if nested_class is not defined
+                        value_dict[key] = nested_class.from_data(data, mapping[key])
+                    else:
+                        # It's a direct reference, get from source data
+                        value_dict[key] = get_recursive(data, mapping[key])
+                else:
+                    # Not a Data subclass, handle normally
+                    value_dict[key] = get_recursive(data, mapping[key])
+        return cls(**value_dict)
+
+
+# Number = float | int

simstack/core/definitions.py

@@ -0,0 +1,21 @@
+from enum import Enum
+
+
+class TaskStatus(str, Enum):
+    SUBMITTED = "submitted"
+    RUNNING = "running"
+    SLURM_QUEUED = "slurm_queued"
+    SLURM_RUNNING = "slurm_running"
+    TIME_OUT = "timeout"
+    FAILED = "failed"
+    COMPLETED = "completed"
+    RECOVERED = "recovered"
+
+
+# TODO: eliminate
+class DBType(str, Enum):
+    SQLITE = "sqlite"
+    MONGODB = "mongodb"
+    POSTGRES = "postgres"
+    IN_MEMORY = "in_memory"
+    WITH_PATH = "with_path"

simstack/core/engine.py

@@ -0,0 +1,116 @@
+from contextvars import ContextVar
+from typing import Optional, Any, Iterable
+
+from odmantic import AIOEngine
+
+
+class AIOEngineProxy(AIOEngine):
+    """
+    A proxy engine that inherits all behavior from AIOEngine but overrides `save`.
+
+    The overridden `save` will:
+    - If the model has a `save` member function, call it with the engine instance.
+    - Else, if any direct part/attribute of the model has a `save` member function, call those with the engine instance.
+    - Otherwise, fall back to the original AIOEngine.save implementation.
+
+    Notes:
+    - Supports single model instances or iterables of model instances (list/tuple/set).
+    - Avoids recursion by delegating to super().save for the fallback.
+    """
+
+    async def save(self, obj: Any, *args, **kwargs) -> Any:  # type: ignore[override]
+        # Handle collections of models
+        if isinstance(obj, (list, tuple, set)):
+            results = []
+            for item in obj:
+                results.append(await self._save_one(item, *args, **kwargs))
+            return results
+        # Single model
+        return await self._save_one(obj, *args, **kwargs)
+
+    async def save_unchecked(self, obj: Any, *args, **kwargs) -> Any:
+        """
+        Save without checking for custom save methods.
+        This is a direct call to the original AIOEngine.save.
+        """
+        return await super().save(obj, *args, **kwargs)
+
+    async def _save_one(self, model: Any, *args, **kwargs) -> Any:
+        # 1) Try model's own `save`
+        if await self._maybe_call_custom_save(model):
+            return None
+
+        # 2) Try parts' `save` (shallow scan of attributes and common containers)
+        parts_saved = await self._call_parts_saves(model)
+
+        # 3) If neither the model nor any part handled saving, fallback to AIOEngine
+        if not parts_saved:
+            return await super().save(model, *args, **kwargs)
+        return None
+
+    async def _maybe_call_custom_save(self, target: Any) -> bool:
+        """
+        If `target` has a callable `save`, invoke it with this engine.
+        Tries positional engine and keyword engine forms, awaits if coroutine.
+        Returns True if a save was invoked, False otherwise.
+        """
+        if not hasattr(target, "save"):
+            return False
+        save_attr = getattr(target, "save")
+        if not callable(save_attr):
+            raise AttributeError(f"Model {target} has no callable `save` method")
+        return await target.save(self)
+
+    async def _call_parts_saves(self, model: Any) -> bool:
+        """
+        Inspect immediate parts/attributes of the model and call their `save` if present.
+        Handles:
+          - Direct attributes from vars(model)
+          - Items inside lists/tuples/sets
+          - Values of dicts
+        Returns True if at least one part save was called.
+        """
+        any_saved = False
+
+        def iter_parts(root: Any) -> Iterable[Any]:
+            # Direct attributes
+            try:
+                for v in vars(root).values():
+                    yield v
+            except TypeError:
+                # Objects without __dict__
+                pass
+
+            # For convenience, also scan common container-typed attributes one level deep
+            try:
+                for v in vars(root).values():
+                    if isinstance(v, (list, tuple, set)):
+                        for item in v:
+                            yield item
+                    elif isinstance(v, dict):
+                        for item in v.values():
+                            yield item
+            except TypeError:
+                pass
+
+        seen_ids = set()
+        for part in iter_parts(model):
+            pid = id(part)
+            if pid in seen_ids:
+                continue
+            seen_ids.add(pid)
+
+            if await self._maybe_call_custom_save(part):
+                any_saved = True
+
+        return any_saved
+
+
+current_engine_context: ContextVar[Optional[AIOEngineProxy]] = ContextVar(
+    "current_engine", default=None
+)
+
+
+def get_current_engine_from_context() -> Optional[AIOEngineProxy]:
+    """Get the current engine from context"""
+    return current_engine_context.get()

simstack/core/find_simstack_modules.py

@@ -0,0 +1,43 @@
+import importlib
+import pkgutil
+from importlib.metadata import entry_points
+from logging import getLogger
+
+logger = getLogger("find_simstack_modules")
+
+
+def discover_simstack_modules():
+    """Discover all packages and modules within the simstack package."""
+
+
+def find_simstack_modules():
+    """Find all packages and modules within the simstack package."""
+
+    all_modules = []
+
+    def walk_packages(package_name):
+        """Walk through all packages and modules."""
+        try:
+            package = importlib.import_module(package_name)
+            package_path = package.__path__
+
+            logger.info(f"Package: {package_name}")
+            # print(f"Path: {package_path}")
+
+            for importer, modname, ispkg in pkgutil.walk_packages(
+                package_path, prefix=f"{package_name}."
+            ):
+                if ispkg:
+                    logger.debug(f"  Subpackage: {modname}")
+                else:
+                    logger.debug(f"  Module: {modname}")
+                    # Split module name by periods
+                    all_modules.append(modname)
+
+        except Exception as e:
+            logger.error(f"Error walking {package_name}: {e}")
+
+    entry_point_list = entry_points(group="simstack.modules")
+    for entry_point in entry_point_list:
+        walk_packages(entry_point.value)
+    return all_modules

simstack/core/hash.py

@@ -0,0 +1,177 @@
+import hashlib
+import inspect
+
+from odmantic import ObjectId
+
+max_iterable_hash_count = 10000
+
+hash_exclusions = [
+    "pandas",
+    "numpy",
+    "ase",
+    "rdkit",
+    "openbabel",
+    "pybel",
+    "rdkit",
+    "pymatgen",
+    "ase",
+    "numpy",
+    "sqlalchemy",
+]
+
+
+def is_iterable(obj):
+    try:
+        iter(obj)
+        return True
+    except TypeError:
+        return False
+
+
+def is_primitive_type(obj):
+    return isinstance(obj, (int, float, bytes, bool, bytearray, type(None)))
+
+
+def hash_value(value):
+    return hashlib.sha256(str(value).encode()).hexdigest()
+
+
+def hash_class_def(cls):
+    try:
+        source_code = inspect.getsource(cls)
+        source_hash = hashlib.sha256(source_code.encode("utf-8")).hexdigest()
+        return source_hash
+    except Exception:
+        mro = cls.__mro__
+        for mro_class in mro:
+            try:
+                source_code = inspect.getsource(mro_class)
+                source_hash = hashlib.sha256(source_code.encode("utf-8")).hexdigest()
+                return source_hash
+            except Exception:
+                pass
+        return "no source code"
+
+
+def hash_iterable(iterable):
+    hash_value = b""
+    for count, item in enumerate(iterable):
+        hash_value += complex_hash_function(item)
+        if count > max_iterable_hash_count:
+            break
+    return hashlib.sha256(hash_value.encode("utf-8")).hexdigest()
+
+
+def hash_function_body(func):
+    # Get the source code of the function
+    source_code = inspect.getsource(func)
+    # Compute the hash of the source code
+    hash_object = hashlib.sha256(source_code.encode("utf-8"))
+    hash_digest = hash_object.hexdigest()
+    return hash_digest
+
+
+def hash_non_callable_members(instance):
+    hashed_values = {}
+    for attr_name, attr_value in vars(instance).items():
+        if not callable(attr_value):
+            hashed_values[attr_name] = hashlib.sha256(
+                str(attr_value).encode()
+            ).hexdigest()
+    return hashed_values
+
+
+class ComplexHash:
+    def __init__(self, obj):
+        self.hash_history = []
+
+    def hash_dict(self, obj):
+        hashed_values = {}
+        for k, v in obj.items():
+            hashed_values[k] = self.complex_hash(v)
+        combined_hash = "".join(f"{k}:{v}" for k, v in sorted(hashed_values.items()))
+        return hashlib.sha256(combined_hash.encode("utf-8")).hexdigest()
+
+    def hash_class(self, cls_obj):
+        # check is the class name starts with a name in the hash_exclusions list
+        class_type = cls_obj.__class__
+        class_name = class_type.__module__ + "." + class_type.__name__
+        if class_type.__module__ == "builtins":
+            return "builtin"
+        if any([class_name.startswith(exclusion) for exclusion in hash_exclusions]):
+            return "excluded"
+        # print("hashing class", cls_obj.__class__.__name__)
+        class_hash = hash_class_def(cls_obj.__class__)
+        # TODO is __dict__ better than vars ?
+        # dict_hash = self.hash_dict(vars(cls_obj))
+
+        obj_dict = cls_obj.__dict__.copy()
+        if hasattr(cls_obj, "__class__") and hasattr(cls_obj.__class__, "__bases__"):
+            for base in cls_obj.__class__.__bases__:
+                if base.__name__ == "Model" and "odmantic.model" in str(
+                    base.__module__
+                ):
+                    if "id" in obj_dict:
+                        del obj_dict["id"]
+
+        dict_hash = self.hash_dict(obj_dict)
+        if hasattr(cls_obj, "model_extra") and cls_obj.model_extra is not None:
+            dict_hash = dict_hash + self.hash_dict(cls_obj.model_extra)
+        if hasattr(cls_obj, "model_fields") and cls_obj.model_fields is not None:
+            dict_hash = dict_hash + self.hash_dict(cls_obj.model_fields)
+        combined_hash = class_hash + dict_hash
+        return hashlib.sha256(combined_hash.encode("utf-8")).hexdigest()
+
+    def complex_hash(self, obj):
+        # TODO: are the functions of a class hashed correctly?
+
+        if isinstance(obj, type):
+            return hash_class_def(obj)
+        elif is_primitive_type(obj):
+            return hash(obj)
+        elif isinstance(obj, ObjectId):
+            return str(obj)
+        elif isinstance(
+            obj, str
+        ):  # strings are iterable but we want to hash them directly
+            return hash_value(obj)
+        elif inspect.isfunction(obj):
+            return hash_function_body(obj)
+        elif inspect.ismethod(obj):
+            if hasattr(obj, "__self__"):
+                if obj in self.hash_history:
+                    return "recursive"
+                self.hash_history.append(obj)
+                if hasattr(obj, "complex_hash"):
+                    return obj.complex_hash()
+                return self.hash_class(obj.__self__)
+            else:
+                return hash_function_body(obj)
+        # this should be instantiated classes
+        elif hasattr(obj, "__dict__"):
+            if obj in self.hash_history:
+                return "recursive"
+            self.hash_history.append(obj)
+            if hasattr(obj, "complex_hash"):
+                return obj.complex_hash()
+            return self.hash_class(obj)
+        elif isinstance(obj, dict):
+            return self.hash_dict(obj)
+        elif is_iterable(obj):
+            hashed_values = {}
+            count = 0
+            for i, item in enumerate(obj):
+                hashed_values[i] = self.complex_hash(item)
+                count += 1
+                if count > max_iterable_hash_count:
+                    break
+            combined_hash = "".join(
+                f"{k}:{v}" for k, v in sorted(hashed_values.items())
+            )
+            return hashlib.sha256(combined_hash.encode("utf-8")).hexdigest()
+        else:
+            return hash_value(obj)
+
+
+def complex_hash_function(obj):
+    return ComplexHash(obj).complex_hash(obj)