simstack 2.0.1__py3-none-any.whl

simstack/tables/model_table.py

@@ -0,0 +1,191 @@
+import inspect
+import json
+import logging
+from pathlib import Path
+
+from odmantic.exceptions import DocumentParsingError
+from pydantic import ValidationError
+
+from simstack.core.context import context
+from simstack.models.models import ModelMapping
+from simstack.models.simstack_model import is_simstack_model
+from simstack.util.db import Database
+from simstack.util.path_manager import path_manager
+from simstack.tables.table_builder import TableBuilderBase
+
+logger = logging.getLogger("ModelTable")
+
+
+class CreateModelTable(TableBuilderBase):
+    """
+    Helper class to build the model table without passing around many parameters.
+
+    Usage:
+        creator = CreateModelTable(db)
+        await creator.make_model_table()
+    """
+
+    @property
+    def logger(self) -> logging.Logger:
+        return logger
+
+    async def build(self, *args, **kwargs) -> None:
+        await super().build(*args, **kwargs)
+        await context.refresh_mappings(models=True, nodes=False)
+
+    async def _process_module(self, module, drops: str) -> None:
+        await self._create_models_from_module(module, drops)
+
+    async def _create_models_from_module(self, module, drops: str):
+        """Create ModelMapping entries for all relevant classes in a module."""
+        classes = inspect.getmembers(module, inspect.isclass)
+
+        for class_name, new_class in classes:
+            # this is required because of the Odmantic Metaclass Model
+            # subclass does not work even if applied to the imported classes
+            # this is a bug in importlib
+            bases = [base.__name__ for base in new_class.__bases__]
+            is_ui_model = any("UIModel" in s for s in bases) or is_simstack_model(
+                new_class
+            )
+            is_model = any(s == "Model" for s in bases)
+
+            is_embedded_model = any(s == "EmbeddedModel" for s in bases)
+
+            if not (is_model or is_ui_model or is_embedded_model):
+                continue
+            if class_name == "Model":
+                continue
+            if new_class.__module__ != module.__name__:
+                continue
+
+            new_modules = new_class.__module__.split(".")
+            # Use drops from the path_info dictionary
+            if drops != "":
+                drop_modules = drops.split(".")
+                while (
+                    new_modules and drop_modules and new_modules[0] == drop_modules[0]
+                ):
+                    new_modules.pop(0)
+                    drop_modules.pop(0)
+                # if len(drop_modules) > 0:
+                #     raise ValueError(
+                #         "drop modules not empty: ", drop_modules, new_class.__module__
+                #     )
+            full_mapping = ".".join(new_modules) + "." + class_name
+            logger.debug(f"    Class: {class_name} Model Mapping: {full_mapping}")
+
+            # Remove any existing ModelMapping entry for this class
+
+            try:
+                existing_entries = await self.db.find(
+                    ModelMapping, ModelMapping.name == class_name
+                )
+            except (ValidationError, NameError, DocumentParsingError) as e:
+                logger.error(
+                    f"Error finding existing ModelMapping for {class_name}: {e}"
+                )
+                existing_entries = []
+
+            if len(existing_entries) > 1:
+                error_msg = f"Fatal error: Found {len(existing_entries)} ModelMapping entries for class_name '{class_name}'. Expected at most one."
+                logger.fatal(error_msg)
+                raise RuntimeError(error_msg)
+
+            if len(existing_entries) == 1:
+                existing_mapping = existing_entries[0].mapping
+                if existing_mapping != full_mapping:
+                    logger.warning(
+                        f"Replacing ModelMapping entry for {class_name}: "
+                        f"old mapping '{existing_mapping}' -> new mapping '{full_mapping}'"
+                    )
+
+                await self.db.delete(existing_entries[0])
+                logger.debug(f"Deleted ModelMapping entry for {class_name}")
+
+            # EmbeddedModels have no collection by may be simstack_models. They are never saved/retrieved
+            collection_name = getattr(new_class, "__collection__", None)
+            if collection_name is None:
+                if is_embedded_model:
+                    collection_name = "EmbeddedModel"
+                else:
+                    logger.error(f"No collection specified for {class_name}")
+
+            # Create the new ModelMapping entry (pickle functionality removed)
+            if is_ui_model:
+                model_entry = ModelMapping(
+                    name=class_name,
+                    mapping=full_mapping,
+                    collection_name=collection_name,
+                    json_schema=json.dumps(new_class.json_schema()),
+                    ui_schema=json.dumps(new_class.ui_schema()),
+                    route="",
+                )
+                logger.debug(
+                    f"SimStack Model: {class_name} Mapping: {full_mapping} Collection: {collection_name}"
+                )
+                # open a file in a subdirectory of the current file schema/model.json
+                if self.write_schema:
+                    project_root = self.project_root
+                    json_file_dir = project_root / "schema"
+                    json_file_dir.mkdir(parents=True, exist_ok=True)
+
+                    combined_schema = {
+                        "json_schema": new_class.json_schema(),
+                        "ui_schema": new_class.ui_schema(),
+                    }
+                    with open(json_file_dir / f"{class_name}.json", "w") as f:
+                        f.write(json.dumps(combined_schema, indent=4))
+            else:
+                model_entry = ModelMapping(
+                    name=class_name,
+                    mapping=full_mapping,
+                    collection_name=collection_name,
+                )
+                logger.debug(
+                    f"Model: {class_name} Mapping: {full_mapping} Collection: {collection_name}"
+                )
+
+            await self.db.save(model_entry)
+
+    async def _make_models_for_path(self, path_name: str):
+        """Build model mappings for all Python files under a configured path."""
+        path_info = path_manager.get_path(path_name)
+        path = path_info["path"]
+        drops = path_info.get("drops", "")
+        logger.info(f"Making model_table entries for files in {path}")
+
+        # Process each file in this path
+        for file_path in path_manager.find_python_files(path_name):
+            await self._create_model_models_from_file(file_path, drops)
+
+    async def clear_table(self) -> None:
+        self.logger.info("Clearing ModelMapping collection")
+        await self.db.get_collection(ModelMapping).drop()
+
+
+# Public API preserved for existing callers (e.g. tests)
+async def make_model_table(
+    db : Database,
+    dirs: list[str] = None,
+    drops: str = "",
+    write_schema: bool = False,
+    clear: bool = False,
+    project_root: Path = None,
+    ignore_entrypoints: bool = False,
+):
+    """
+    Rebuild the model table using the given database.
+
+    This is a thin wrapper around CreateModelTable for backward compatibility.
+    """
+    creator = CreateModelTable(db, write_schema=write_schema, project_root=project_root)
+    await creator.build(dirs=dirs, drops=drops, clear=clear, ignore_entrypoints=ignore_entrypoints)
+
+
+def create_model_table_main():
+    TableBuilderBase.cli_main(CreateModelTable)
+
+
+if __name__ == "__main__":
+    create_model_table_main()

simstack/tables/node_children.py

@@ -0,0 +1,219 @@
+import ast
+import importlib
+import inspect
+import os
+from typing import Optional, Callable, List, Dict
+from simstack.models import NodeModel
+from simstack.util.importer import import_function
+from simstack.util.docstring_parser import DocstringParser
+
+import logging
+logger = logging.getLogger("node_children")
+
+def _import_callable_from_mapping(function_mapping: str) -> Optional[Callable]:
+    """
+    Import a callable from a 'module.submodule.func' mapping string.
+    """
+    try:
+        module_path, _, attr = function_mapping.rpartition(".")
+        if not module_path or not attr:
+            logger.warning(f"Invalid function_mapping '{function_mapping}' (expected 'module.func').")
+            return None
+
+        module = importlib.import_module(module_path)
+        obj = getattr(module, attr, None)
+        if obj is None or not callable(obj):
+            logger.warning(f"Could not resolve callable for '{function_mapping}'.")
+            return None
+
+        return obj
+    except Exception as e:
+        logger.warning(f"Could not import '{function_mapping}': {e}")
+        return None
+
+def _extract_called_functions(func: Callable) -> List[str]:
+    """
+    Inspect the function code and return all the function calls made within it.
+
+    Returns:
+        List of function names called within the function.
+    """
+    called_functions: List[str] = []
+
+    try:
+        source = inspect.getsource(func)
+        tree = ast.parse(source)
+
+        task_creators = {"create_tasks", "ensure_future"}
+
+        for node in ast.walk(tree):
+            if isinstance(node, ast.Call):
+                if isinstance(node.func, ast.Name):
+                    # Direct function call: func_name()
+                    called_functions.append(node.func.id)
+                elif isinstance(node.func, ast.Attribute):
+                    # Attribute function call: obj.method()
+                    if isinstance(node.func.value, ast.Name):
+                        called_functions.append(f"{node.func.value.id}.{node.func.attr}")
+                    else:
+                        called_functions.append(node.func.attr)
+
+                # Detect async task creation: create_tasks(func()), ensure_future(func()), etc.
+                if isinstance(node.func, ast.Name) and node.func.id in task_creators:
+                    if node.args and isinstance(node.args[0], ast.Call):
+                        task_call = node.args[0]
+                        if isinstance(task_call.func, ast.Name):
+                            called_functions.append(task_call.func.id)
+                        elif isinstance(task_call.func, ast.Attribute):
+                            if isinstance(task_call.func.value, ast.Name):
+                                called_functions.append(f"{task_call.func.value.id}.{task_call.func.attr}")
+                            else:
+                                called_functions.append(task_call.func.attr)
+
+                elif isinstance(node.func, ast.Attribute) and node.func.attr in task_creators:
+                    if node.args and isinstance(node.args[0], ast.Call):
+                        task_call = node.args[0]
+                        if isinstance(task_call.func, ast.Name):
+                            called_functions.append(task_call.func.id)
+                        elif isinstance(task_call.func, ast.Attribute):
+                            if isinstance(task_call.func.value, ast.Name):
+                                called_functions.append(f"{task_call.func.value.id}.{task_call.func.attr}")
+                            else:
+                                called_functions.append(task_call.func.attr)
+
+    except Exception as e:
+        logger.warning(f"Could not extract called functions from {getattr(func, '__name__', '<unknown>')}: {e}")
+
+    return called_functions
+
+def normalize_text(s: str) -> str:
+    return s.encode("utf-8", errors="replace").decode("utf-8")
+
+async def update_node_children(database, drops: str) -> None:
+    """
+    1) Read all registered nodes from NodeModel.
+    2) For each node, import its function and extract called functions.
+    3) Keep only those calls that correspond to nodes in the table.
+    4) Update the NodeModel.called_nodes of the original function with a list of *function_mapping* strings.
+    """
+    node_models: List[NodeModel] = await database.find(NodeModel)
+
+    # Build lookup tables to resolve extracted names -> NodeModel.function_mapping
+    mapping_by_node_name: Dict[str, str] = { nm.name: nm.function_mapping for nm in node_models}
+    mapping_set: set[str] = set(nm.function_mapping for nm in node_models)
+
+    for nm in node_models:
+        func = await import_function(nm.function_mapping, database, None, True)
+        if func is None:
+            continue
+
+        parser = DocstringParser(inspect.getdoc(func))
+
+        called = _extract_called_functions(func)
+        logger.debug(f"Node {nm.name} called: {called}")
+        called_from_docstring = parser.called_nodes()
+        resolved: set[str] = set()
+        for called_name in called:
+            # If the extractor ever returns a full mapping, accept it directly
+            if called_name in mapping_set:
+                resolved.add(called_name)
+                continue
+
+            # If it matches a node name, resolve via node name
+            if called_name in mapping_by_node_name:
+                resolved.add(mapping_by_node_name[called_name])
+                continue
+
+            # Otherwise, resolve by last segment (handles "obj.method" and "module.func")
+            short = called_name.split(".")[-1]
+            if short in mapping_by_node_name:
+                resolved.add(mapping_by_node_name[short])
+                if called_from_docstring is not None and short in called_from_docstring:
+                    del called_from_docstring[short]
+                continue
+
+        # go over the leftovers in the docstring
+        if called_from_docstring is not None:
+            for called_name in called_from_docstring:
+                if called_name in mapping_set:
+                    resolved.add(called_name)
+                    continue
+
+                # If it matches a node name, resolve via node name
+                if called_name in mapping_by_node_name:
+                    resolved.add(mapping_by_node_name[called_name])
+                    continue
+
+                # Otherwise, resolve by last segment (handles "obj.method" and "module.func")
+                # short = called_name.split(".")[-1]
+                # if short in mapping_by_node_name:
+                #     resolved.add(mapping_by_short_name[short])
+                #     continue
+
+        if len(resolved) != 0:
+            logger.debug(f"Resolved children of {nm.name} to {resolved}")
+        nm.called_nodes = sorted(resolved)
+
+        
+        await database.get_collection(NodeModel).update_one(
+            {"_id": nm.id},
+            {"$set": {"called_nodes": nm.called_nodes}},
+        )
+
+    with open("node_models.txt", "w") as outfile:
+        for nm in node_models:
+            mapping_set.add(nm.function_mapping)
+            mapping_by_node_name[nm.name] = nm.function_mapping
+
+            # Write formatted NodeModel information
+            outfile.write(f"{'=' * 80}\n")
+            
+            outfile.write(f"Node: {nm.name if nm.name else 'WARNING: Missing name'}\n")
+            outfile.write(
+                f"Function Mapping: {nm.function_mapping if nm.function_mapping else 'WARNING: Missing function_mapping'}\n")
+
+            if hasattr(nm, 'description') and nm.description:
+                try:
+                    cleaned_description = normalize_text(nm.description)
+                    outfile.write(f"Description: {cleaned_description}\n")
+                except Exception as e:
+                    logger.warning(f"Could not normalize description for {nm.name}: {e}")
+            else:
+
+                outfile.write(f"Description: WARNING: Missing description\n")
+
+            # Write input_mappings in table format
+            outfile.write(f"\nInput Mappings:\n")
+            if hasattr(nm, 'input_mappings') and nm.input_mappings:
+                outfile.write(f"  {'Arg Name':<20} | {'Model':<50} | {'Field':<20}\n")
+                outfile.write(f"  {'-' * 20}-+-{'-' * 50}-+-{'-' * 20}\n")
+                for mapping in nm.input_mappings:
+                    arg_name = getattr(mapping, 'name', 'N/A')
+                    model = getattr(mapping, 'mapping', 'N/A')
+                    field = getattr(mapping, 'description', 'N/A')
+                    outfile.write(f"  {str(arg_name):<20} | {str(model):<50} | {str(field):<20}\n")
+            else:
+                outfile.write(f"  (none)\n")
+
+            # Write result_mappings in table format
+            outfile.write(f"\nResult Mappings:\n")
+            if hasattr(nm, 'result_mappings') and nm.result_mappings:
+                outfile.write(f"  {'Arg Name':<20} | {'Model':<50} | {'Field':<20}\n")
+                outfile.write(f"  {'-' * 20}-+-{'-' * 50}-+-{'-' * 20}\n")
+                for mapping in nm.result_mappings:
+                    arg_name = getattr(mapping, 'name', 'N/A')
+                    model = getattr(mapping, 'mapping', 'N/A')
+                    field = getattr(mapping, 'description', 'N/A')
+                    outfile.write(f"  {str(arg_name):<20} | {str(model):<50} | {str(field):<20}\n")
+            else:
+                outfile.write(f"  (none)\n")
+
+            # Write called_nodes (node children)
+            outfile.write(f"\nCalled Nodes (Children):\n")
+            if hasattr(nm, 'called_nodes') and nm.called_nodes:
+                for called_node in nm.called_nodes:
+                    outfile.write(f"  - {called_node}\n")
+            else:
+                outfile.write(f"  (none)\n")
+
+            outfile.write(f"{'-' * 80}\n\n")