simstack 2.0.1__py3-none-any.whl

simstack/tables/table_builder.py

@@ -0,0 +1,314 @@
+import argparse
+import asyncio
+import importlib
+import logging
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Iterable, Optional, Type
+import fnmatch  # <-- added
+
+from simstack.core.context import context
+from simstack.core.find_simstack_modules import find_simstack_modules
+from simstack.util.db import Database
+from simstack.util.import_module import import_module_from_file
+from simstack.util.path_manager import path_manager
+
+
+
+class TableBuilderBase(ABC):
+    """
+    Shared pipeline for building "tables" by scanning:
+      1) installed simstack modules (packages)
+      2) configured project paths (python files)
+
+    Subclasses only implement `_process_module(module, drops)`.
+    """
+
+    def __init__(self, db: Database, write_schema: bool = False, project_root: Optional[Path] = None):
+        self.db = db
+        self.write_schema = write_schema
+        self._project_root = project_root
+        self._processed_modules = set()
+
+    @property
+    def project_root(self) -> Path:
+        if self._project_root:
+            return self._project_root
+        if context.config:
+            return context.config.project_root
+        from simstack.util.project_root_finder import find_project_root
+        return find_project_root()
+
+    @property
+    @abstractmethod
+    def logger(self) -> logging.Logger:
+        raise NotImplementedError
+
+    async def build(
+        self,
+        *,
+        dirs: Optional[list[Path]] = None,
+        drops: str = "",
+        exclude: Optional[list[str]] = None,
+        clear: bool = False,
+        ignore_entrypoints: bool = False,
+    ) -> None:
+        """
+        Build the table.
+
+        - If `dirs` is None: use configured `path_manager` paths (existing behavior).
+        - If `dirs` is a list: scan those directories for Python files and process them.
+
+        `exclude` entries can match:
+          - path parts (e.g. ".venv", "__pycache__")
+          - glob patterns (e.g. "*.generated.py")
+          - nested relative paths (e.g. "src/simstack/models")
+        """
+        await self._ensure_context_initialized()
+        self._processed_modules.clear()
+        if clear:
+            await self.clear_table()
+
+        if not ignore_entrypoints:
+            await self._process_simstack_modules(drops=drops)
+
+        await self._process_dirs(dirs, drops=drops, exclude=exclude or [])
+
+    async def _ensure_context_initialized(self) -> None:
+        if not context.initialized:
+            await context.initialize()
+
+    async def _process_simstack_modules(self, drops: str) -> None:
+        all_modules = set(find_simstack_modules())
+        for module_name in all_modules:
+            if module_name in self._processed_modules:
+                continue
+            self.logger.debug("Processing module: %s", module_name)
+            module = self._import_package_module(module_name)
+            if module is None:
+                continue
+            await self._process_module(module, drops=drops)
+            self._processed_modules.add(module_name)
+
+    def _import_package_module(self, module_name: str):
+        try:
+            return importlib.import_module(module_name)
+        except Exception as exc:
+            self.logger.warning("Failed to import module %s: %s", module_name, exc)
+            return None
+
+    async def _process_configured_paths(self) -> None:
+        for path_name in path_manager.paths.keys():
+            await self._process_path(path_name)
+
+    async def _process_path(self, path_name: str) -> None:
+        path_info = path_manager.get_path(path_name)
+        base_path = path_info["path"]
+        drops = path_info.get("drops", "")
+
+        self.logger.info("Processing configured path %s: %s", path_name, base_path)
+
+        for file_path in path_manager.find_python_files(path_name):
+            await self._process_file(file_path, drops)
+
+    async def _process_dirs(self, dirs: list[Path], *, drops: str, exclude: list[str]) -> None:
+        for base_dir in dirs:
+            self.logger.info("Processing CLI dir: %s", base_dir)
+
+            base_dir = Path(base_dir)
+
+            # Accept either absolute paths or paths relative to project root.
+            base_dir_path = base_dir if base_dir.is_absolute() else (self.project_root / base_dir)
+
+            for py_file in self._iter_python_files_under_dir(base_dir_path, exclude=exclude):
+                await self._process_file(py_file, drops)
+
+    def _iter_python_files_under_dir(self, base_dir: Path, *, exclude: list[str]) -> Iterable[Path]:
+        default_exclude_parts = {
+            ".git",
+            ".hg",
+            ".svn",
+            ".venv",
+            "venv",
+            "__init__.py"
+            "__pycache__",
+            ".mypy_cache",
+            ".pytest_cache",
+            ".ipynb_checkpoints"
+        }
+
+        def _should_exclude(p: Path) -> bool:
+            # Fast path: ignore common tooling/cache dirs anywhere in the path.
+            if any(part in default_exclude_parts for part in p.parts):
+                return True
+
+            # If the caller didn't pass any excludes, we're done.
+            if not exclude:
+                return False
+
+            # Match excludes against path parts and relative path (for nested patterns).
+            try:
+                rel = p.relative_to(base_dir)
+            except ValueError:
+                rel = p
+
+            rel_posix = rel.as_posix()
+
+            for ex in exclude:
+                ex = (ex or "").strip()
+                if not ex:
+                    continue
+
+                # 1) Exact match against any path segment (works for "build", ".tox", etc.)
+                if ex in p.parts or ex in rel.parts:
+                    return True
+
+                # 2) Glob match against the relative posix path (supports nested paths + file globs)
+                #    Examples:
+                #      --exclude "src/simstack/models"
+                #      --exclude "**/generated/**"
+                #      --exclude "*.generated.py"
+                if fnmatch.fnmatch(rel_posix, ex) or fnmatch.fnmatch(rel_posix, ex.rstrip("/") + "/*"):
+                    return True
+
+                # 3) Also allow Windows-ish inputs like "a\\b\\c" by normalizing to posix.
+                ex_posix = ex.replace("\\", "/")
+                if ex_posix != ex and (
+                    fnmatch.fnmatch(rel_posix, ex_posix) or fnmatch.fnmatch(rel_posix, ex_posix.rstrip("/") + "/*")
+                ):
+                    return True
+
+            return False
+
+        if base_dir.is_file() and base_dir.suffix == ".py":
+            if not _should_exclude(base_dir):
+                yield base_dir
+            return
+
+        if not base_dir.exists():
+            self.logger.warning("Skipping non-existent path: %s", base_dir)
+            return
+
+        if not base_dir.is_dir():
+            self.logger.warning("Skipping non-directory path: %s", base_dir)
+            return
+
+        for p in base_dir.rglob("*.py"):
+            if _should_exclude(p):
+                continue
+            if p.is_file():
+                yield p
+
+    async def _process_file(self, file_path: Path, drops: str) -> None:
+        self.logger.debug("Processing file: %s", file_path)
+        module = import_module_from_file(file_path, self.project_root)
+        if not module:
+            self.logger.debug("Skipping %s because module import returned None", file_path)
+            return
+        if module.__name__ in self._processed_modules:
+            self.logger.debug("Skipping %s because module %s was already processed", file_path, module.__name__)
+            return
+        await self._process_module(module, drops)
+        self._processed_modules.add(module.__name__)
+
+    @abstractmethod
+    async def _process_module(self, module, drops: str) -> None:
+        """Subclass hook: scan/register whatever you need from `module`."""
+        raise NotImplementedError
+
+    async def second_stage(self, drops: str) -> None:
+        """ an optional hook to run after all modules have been processed"""
+        pass
+
+    async def clear_table(self) -> None:
+        """ an optional hook to clear the table before building """
+        pass
+
+    @classmethod
+    def cli_main(cls, builder_cls: Type["TableBuilderBase"], write_schema: bool = False) -> None:
+        """
+        Reusable CLI entry point.
+
+        Options:
+          --dir RELPATH   (repeatable) Directories (relative to CWD) to scan for *.py files.
+                         If omitted, defaults to scanning the current working directory.
+          --drops STRING  Drops prefix (string) applied while processing these CLI dirs.
+          -v/--verbose    Increase logging verbosity.
+        """
+        parser = argparse.ArgumentParser()
+        parser.add_argument("-v", "--verbose", action="count", default=0)
+        parser.add_argument(
+            "--dir",
+            dest="dirs",
+            action="append",
+            default=[],
+            help="Path to scan (repeatable). If omitted, CWD is used.",
+        )
+        parser.add_argument(
+            "--exclude",
+            dest="exclude",
+            action="append",
+            default=[],
+            help=(
+                "Exclude a directory/name/glob from scanning (repeatable). "
+          
+                "Examples: --exclude .venv --exclude __pycache__ --exclude src/simstack/models --exclude '*.generated.py'"
+            ),
+        )
+        parser.add_argument(
+            "--drops",
+            dest="drops",
+            default="",
+            help="Drops prefix used when processing CLI dirs (string).",
+        )
+        parser.add_argument(
+            "--write-schema",
+            dest="write_schema",
+            action="store_true",
+            help="Enable schema writing.",
+        )
+        parser.add_argument(
+            "--ignore-entrypoints",
+            dest="ignore_entrypoints",
+            action="store_true",
+            help="Do not scan for modules registered via entrypoints.",
+        )
+        parser.add_argument(
+            "--clear",
+            dest="clear",
+            action="store_true",
+            help="Clear the table before building.",
+        )
+
+        args = parser.parse_args()
+
+        level = logging.WARNING
+        if args.verbose == 1:
+            level = logging.INFO
+        elif args.verbose >= 2:
+            level = logging.DEBUG
+
+        # Resolve dirs relative to the current working directory
+        if args.dirs:
+            dirs = [Path.cwd() / d for d in args.dirs]
+        else:
+            dirs = [Path.cwd()]
+
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+
+        logging.getLogger("pymongo").setLevel(logging.INFO)
+
+        async def _run() -> None:
+            await context.initialize(log_level=level, resource="self")
+            builder = builder_cls(context.db, write_schema=args.write_schema)
+            await builder.build(
+                dirs=dirs,
+                drops=args.drops,
+                exclude=args.exclude,
+                clear=args.clear,
+                ignore_entrypoints=args.ignore_entrypoints,
+            )
+            await builder.second_stage(args.drops)
+        loop.run_until_complete(_run())
+        loop.close()

simstack/util/async_zip_utils.py

@@ -0,0 +1,18 @@
+async def async_zip(async_gen1, async_gen2):
+    """Async version of zip for two async generators"""
+    async for item1, item2 in async_zip_multiple(async_gen1, async_gen2):
+        yield item1, item2
+
+
+async def async_zip_multiple(*async_generators):
+    """Async version of zip for multiple async generators"""
+    iterators = [aiter(ag) for ag in async_generators]
+
+    while True:
+        try:
+            values = []
+            for it in iterators:
+                values.append(await anext(it))
+            yield tuple(values)
+        except StopAsyncIteration:
+            break

simstack/util/b64mixin.py

@@ -0,0 +1,23 @@
+import base64
+import zlib
+
+
+class BytesB64Mixin:
+    """
+    Mixin that teaches Pydantic/ODMantic to serialise *bytes* fields as
+    base-64-encoded ASCII strings when exporting to JSON (dict / response).
+    """
+
+    model_config = {
+        "json_encoders": {bytes: lambda b: base64.b64encode(b).decode("ascii")}
+    }
+
+    def _compress_bytes(self, data: bytes) -> str:
+            """Compress bytes and encode to base64 string"""
+            compressed = zlib.compress(data)
+            return base64.b64encode(compressed).decode('utf-8')
+
+    def _decompress_bytes(self, data: str) -> bytes:
+            """Decode base64 string and decompress to bytes"""
+            compressed = base64.b64decode(data.encode('utf-8'))
+            return zlib.decompress(compressed)

simstack/util/cleaned_json_schema.py

@@ -0,0 +1,20 @@
+def cleaned_json_schema(cls):
+    """
+
+    Generates a JSON schema for the given class and its fields, but eliminates
+    all fields which are models, embedded models, or references to models.
+
+    :param cls:
+    :return:
+    """
+    schema = cls.model_json_schema()
+    schema["title"] = cls.__name__
+
+    # for field_name, field in cls.model_fields.items():
+    #     field_type = field.annotation
+    #     if hasattr(field_type, '_is_simstack_model'):
+    #         # If the field is a model or embedded model, remove its proprties
+    #         if field_name in schema['$defs'] and "properties" in schema['$defs']['properties']:
+    #             del schema['$defs']['properties']
+
+    return schema

simstack/util/config_reader.py

@@ -0,0 +1,179 @@
+from pathlib import Path
+from typing import List, TYPE_CHECKING, Union
+from simstack.core.resources import allowed_resources
+from simstack.models.parameters import Resource
+from simstack.models.resource_definition import ResourceDefinition, GitRepo
+from simstack.util.transform_file_name import transform_file_name
+from simstack.util.init_data_source import initialize_resource_from_db
+
+from simstack.util.toml_reader import TomlReader
+from simstack.util.database_information import DatabaseInformation
+
+if TYPE_CHECKING:
+    from simstack.util.db import Database
+
+
+class ConfigReader(DatabaseInformation):
+    """
+    Represents a configuration reader that integrates with a database and a resource
+    definition system.
+    """
+
+    def __init__(
+        self,
+        db_info: Union[DatabaseInformation, "Database"],
+        resource_definition: ResourceDefinition,
+        *,
+        project_root: Path,
+        git_list: list[GitRepo] | None = None,
+    ):
+        db_temp_info = DatabaseInformation.from_db_info_or_db(db_info)
+        super().__init__(*db_temp_info.get_information())
+
+        self._project_root = project_root
+        self._resource_definition = resource_definition
+        self._git_list = list(git_list or [])
+        self._resource_str = resource_definition.resource_str
+
+    @classmethod
+    async def create(
+        cls,
+        resource_str,
+        db: "Database",
+        toml_reader: TomlReader,
+        project_root: Path,
+        **kwargs,
+    ):
+        """
+
+        the idea here is to attempt to get:
+
+        1) all required data from the kwargs, if available
+        2) if the data is not all available, get it from the database
+
+        """
+        import logging
+
+        logger = logging.getLogger("config-reader")  # do this here because the calling function sets the logger up
+
+        required_keys = [
+            "python_paths",
+            "workdir",
+            "environment_start",
+        ]
+        init_done = True
+        workdir_self = None
+        config = {"project_root": project_root}
+        for key in required_keys:
+            if key in kwargs:
+                config[key] = kwargs.get(key)
+                if key == "workdir":
+                    workdir_self = kwargs.get("workdir")
+                logger.info(f"Init from kwargs: {key}: {kwargs.get(key)}")
+            else:
+                init_done = False
+        for key in ["server_url", "server_token"]:  # these are optional
+            if key in kwargs:
+                config[key] = kwargs.get(key)
+                logger.info(f"Init from kwargs: {key}: {kwargs.get(key)}")
+
+        git_list = []
+        if not init_done:
+            if toml_reader is None:
+                raise ValueError("ConfigReader: No toml_reader provided.")
+            workdir_self = kwargs.get("workdir", None)
+            if workdir_self is None:
+                workdir_self = toml_reader.get("parameters.general.workdir_self", None)
+                if workdir_self is None:
+                    workdir_self = toml_reader.get("resources.self.workdir", None)
+            if workdir_self is None:
+                # raise ValueError("No workdir for self specified in config file or keyword arguments.")
+                workdir_self = project_root / "simstack" # Default to project_root / simstack if not found
+            else:
+                workdir_self = Path(workdir_self)
+
+        resource_definition = await initialize_resource_from_db(resource_str, db, workdir_self)
+            # else:
+            #     raise RuntimeError("initialization from toml file is deprecated")
+            #     allowed_resources_list = toml_reader.get_allowed_resources()
+            #     allowed_resources.set_resources(allowed_resources_list)
+            #     resource_definition = toml_reader.get_resource_definition(resource_str)
+            #     git_list = toml_reader.get_git_list()
+            #     toml_reader.build_routes()
+
+        if resource_definition is None:
+            raise ValueError("No valid resource definition found.")
+
+        # override the values in resource definition with those from the keyword arguments
+        for key in resource_definition.__class__.model_fields.keys():
+            if key in config:
+                logger.info(f"Overriding {key} from kwargs to: {config[key]}")
+                resource_definition.__setattr__(key, config[key])
+                del config[key]
+                del required_keys[required_keys.index(key)]
+
+        project_root = config.pop("project_root")
+        git_list_final = config.pop("git_list", git_list)
+
+        if config:
+            logger.warning(
+                f"Ignoring unused ConfigReader init keys: {sorted(config.keys())}"
+            )
+
+        log_msg = f"Resource: {resource_definition.resource_str}"
+        for key, value in resource_definition.__dict__.items():
+            if key in ["workdir", "git-branch", "environment_start"]:
+                log_msg += f" {key}: {value}"
+        logger.info(log_msg)
+
+        return cls(
+            db, resource_definition, project_root=project_root, git_list=git_list_final
+        )
+
+    @property
+    def server_url(self) -> str:
+        return self._server_url
+
+    @property
+    def server_token(self) -> str:
+        return self._server_token
+
+    @property
+    def git_list(self) -> List[GitRepo]:
+        return self._git_list
+
+    @property
+    def resource(self) -> Resource:
+        return Resource(value=self._resource_str)
+
+    @property
+    def workdir(self) -> Path:
+        return transform_file_name(self._resource_definition.workdir, self.project_root)
+
+    @property
+    def environment_start(self) -> str:
+        return self._resource_definition.environment_start
+
+    @property
+    def python_paths(self) -> List[Path]:
+        return [transform_file_name(p) for p in self._resource_definition.python_paths]
+
+    @property
+    def ssh_key(self) -> Path:
+        return transform_file_name(self._resource_definition.ssh_key)
+
+    @property
+    def hostname(self):
+        return self._resource_definition.hostname
+
+    @resource.setter
+    def resource(self, value: str):
+        raise ValueError("ConfigReader: Resource cannot be set directly")
+
+    @property
+    def project_root(self) -> Path:
+        return self._project_root
+
+    @property
+    def docker(self) -> bool:
+        return False

simstack/util/custom_model_dump.py

@@ -0,0 +1,75 @@
+from typing import Dict, Any
+
+
+async def custom_model_dump(self, **kwargs) -> Dict[str, Any]:
+    """
+    Custom model dump method to handle the conversion of model instances to dictionaries.
+    This method recursively traverses dictionaries and lists to convert any nested
+    model instances to their dictionary representation.
+
+    :param self: The model instance
+    :param kwargs: Additional keyword arguments
+    :return: A dictionary representation of the model instance
+    """
+    # First get the dictionary from the standard model_dump
+
+    filtered_kwargs = {k: v for k, v in kwargs.items() if k != "engine"}
+    data = self.model_dump(**filtered_kwargs)
+
+    # Process each field - the key is to process the original values from self
+    for field_name, field in self.model_fields.items():
+        field_value = getattr(self, field_name)
+
+        # Skip None values
+        if field_value is None:
+            continue
+
+        data[field_name] = await _recursive_process(field_value, **kwargs)
+
+    return data
+
+
+async def _recursive_process(obj, **kwargs):
+    """
+    Helper function to recursively process an object that might contain models
+    """
+    if hasattr(obj, "custom_model_dump"):
+        return await obj.custom_model_dump(**kwargs)
+    elif isinstance(obj, dict):
+        return {k: await _recursive_process(v, **kwargs) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [await _recursive_process(item, **kwargs) for item in obj]
+    else:
+        return obj
+
+
+async def _process_container(container, **kwargs):
+    """
+    Helper function to recursively process dictionaries and lists,
+    converting any model instances to their dictionary representation.
+
+    :param container: The container to process (dict or list)
+    :param kwargs: Additional keyword arguments
+    :return: Processed container with model instances converted to dictionaries
+    """
+    if isinstance(container, dict):
+        result = {}
+        for k, v in container.items():
+            if hasattr(v, "custom_model_dump"):
+                result[k] = await v.custom_model_dump(**kwargs)
+            elif isinstance(v, (dict, list)):
+                result[k] = await _process_container(v, **kwargs)
+            else:
+                result[k] = v
+        return result
+    elif isinstance(container, list):
+        result = []
+        for item in container:
+            if hasattr(item, "custom_model_dump"):
+                result.append(await item.custom_model_dump(**kwargs))
+            elif isinstance(item, (dict, list)):
+                result.append(await _process_container(item, **kwargs))
+            else:
+                result.append(item)
+        return result
+    return container