simstack 2.0.1__py3-none-any.whl

simstack/methods/upload_helpers.py

@@ -0,0 +1,81 @@
+import zipfile
+import tarfile
+import os
+from pathlib import Path
+
+from simstack.core.context import context
+from simstack.core.node import node
+from simstack.models import FileListModel
+from simstack.models.files import FileStack
+
+
+@node
+async def archive_upload(file_upload: FileStack, **kwargs):
+    """
+    Asynchronously processes a file upload, extracts files from supported archive formats
+    (.zip, .tar, .tar.gz), and returns a list of extracted files wrapped in FileStack
+    instances. If the file is not an archive, it is added as-is to the output.
+
+    Args:
+        file_upload (FileStack): The uploaded file to be processed.
+        **kwargs: Additional keyword arguments.
+            node_runner: An optional object for logging or processing context. If supplied,
+            it will log file paths and extraction details using its `info` method.
+
+    Returns:
+        FileListModel: A model containing a list of FileStack instances representing
+        the extracted files or the original file if it is not an archive.
+    """
+    node_runner = kwargs.get('node_runner', None)
+    local_file = file_upload.get()
+
+    local_file_path = Path(local_file)
+    node_runner.info(f"Local file path: {local_file_path}")
+    extract_dir = local_file_path.parent
+
+    output_file_list = FileListModel(field_name="extracted_files")
+    if local_file_path.suffix == '.zip':
+        with zipfile.ZipFile(local_file_path, 'r') as zip_ref:
+            node_runner.info(f"Extracting files from {local_file_path} to {extract_dir}")
+            zip_ref.extractall(extract_dir)
+
+        for root, dirs, files in os.walk(extract_dir):
+            for file in files:
+                file_path = Path(root) / file
+                if file_path != local_file_path:
+                    node_runner.info(f"Storing extracted file: {file_path}")
+                    file_stack = FileStack.from_local_file(file_path, in_memory=True, is_hashable=True, secure_source=True)
+                    output_file_list.append(file_stack)
+
+    elif local_file_path.suffix == '.tar' or local_file_path.name.endswith('.tar.gz'):
+        with tarfile.open(local_file_path, 'r:*') as tar_ref:
+            node_runner.info(f"Extracting files from {local_file_path} to {extract_dir}")
+            tar_ref.extractall(extract_dir)
+
+        for root, dirs, files in os.walk(extract_dir):
+            for file in files:
+                file_path = Path(root) / file
+                if file_path != local_file_path:
+                    node_runner.info(f"Storing extracted file: {file_path}")
+                    file_stack = FileStack.from_local_file(file_path, in_memory=True, is_hashable=True, secure_source=True)
+                    output_file_list.append(file_stack)
+    else:
+        output_file_list.append(file_upload)
+
+    return output_file_list
+
+@node
+async def file_list_upload_test(file_list: FileListModel,**kwargs):
+    node_runner = kwargs.get('node_runner', None)
+    # Ensure cache is loaded before iterating synchronously
+    for file_stack in file_list:
+        local_file = file_stack.get()
+        node_runner.info(f"Local file path: {local_file}")
+    return True
+
+@node
+async def file_list_upload_all(file_stack: FileStack,**kwargs):
+    file_list = await archive_upload(file_stack, **kwargs)
+    unpacked = await file_list_upload_test(file_list, **kwargs)
+    return unpacked
+

simstack/models/__init__.py

@@ -0,0 +1,74 @@
+from .fire_and_forget_result import FireAndForgetResult
+from .simstack_model import simstack_model
+from .artifact_models import ArtifactMapping, ArtifactModel
+from .base_types import (
+    IntData,
+    FloatData,
+    StringData,
+    BooleanData,
+    BinaryOperationInput,
+    IteratorInput,
+)
+from .base_lists import StringDataList, StringList
+from .array_list import ArrayList
+from .file_list import FileList, FileListModel
+from .models import (
+    ModelMapping,
+    NodeModel,
+)
+from .files import FileStack
+from .parameters import Parameters
+from .node_registry import NodeRegistry
+from .named_data_reference import NamedDataReference
+from .project import Project
+from .tag import Tag
+from .resource_assignment import ResourceAssignmentRule, SlurmParametersPatch
+from .datasettuple import (
+    DataSetTuple,
+    DataSetTupleSection,
+    DataSetTupleSelection,
+    DataSetTupleSelectionField,
+)
+from .dataset import DataSet, DataSetSection, DataSetSelection, DataSetSelectionField
+from .dataset_metadata import DataSetMetadata, DataSetMetadataTemplate
+from .images2d import Image2DArtifactModel
+
+
+__all__ = [
+    "FireAndForgetResult",
+    "ModelMapping",
+    "NodeModel",
+    "IntData",
+    "FloatData",
+    "StringData",
+    "BooleanData",
+    "ArrayList",
+    "BinaryOperationInput",
+    "IteratorInput",
+    "ArtifactMapping",
+    "ArtifactModel",
+    "FileListModel",
+    "Parameters",
+    "NodeRegistry",
+    "NamedDataReference",
+    "FileStack",
+    "Project",
+    "Tag",
+    "ResourceAssignmentRule",
+    "SlurmParametersPatch",
+    "simstack_model",
+    "DataSet",
+    "DataSetSection",
+    "DataSetSelection",
+    "DataSetTuple",
+    "DataSetTupleSection",
+    "DataSetTupleSelectionField",
+    "DataSetMetadata",
+    "DataSetSelectionField",
+    "DataSetTupleSelection",
+    "FileList",
+    "DataSetMetadataTemplate",
+    "Image2DArtifactModel",
+    "StringDataList",
+    "StringList",
+]

simstack/models/array_list.py

@@ -0,0 +1,21 @@
+from typing import List, Optional
+
+from odmantic import Model, Field, ObjectId
+
+from simstack.models import simstack_model
+from simstack.models.array_storage import ArrayStorage
+from simstack.util.object_list_mixin import ObjectListMixin
+
+@simstack_model
+class ArrayList(Model, ObjectListMixin[ArrayStorage]):
+    elements: List[ObjectId] = Field(default_factory=list)
+
+    def __init__(self, **data):
+        data, cache = self._normalize_elements_for_init(data)
+        Model.__init__(self, **data)
+        if cache is not None:
+            self._set_cache(cache)
+
+    def __iter__(self):
+        return ObjectListMixin.__iter__(self)
+

simstack/models/array_storage.py

@@ -0,0 +1,83 @@
+import json
+import zlib
+from typing import Optional
+
+from odmantic import Model
+from pydantic import model_validator
+
+from simstack.models.simstack_model import simstack_model
+from simstack.util.ui_tools import ui_hide_fields
+from simstack.util.b64mixin import BytesB64Mixin
+
+
+@simstack_model
+class ArrayStorage(BytesB64Mixin, Model):
+    name: Optional[str]
+    shape: Optional[str] = None  # Store array shape as string like "3,3"
+    field_name: Optional[str] = None  # Store flattened array data as compressed JSON
+    data_json: Optional[str] = None
+
+    @model_validator(mode='before')
+    @classmethod
+    def copy_name_to_field_name(cls, values):
+        if isinstance(values, dict) and 'name' in values and 'field_name' not in values:
+            values['field_name'] = values['name']
+        return values
+
+    def set_array(self, array):
+        """Store a numpy array"""
+        self.shape = ",".join(str(dim) for dim in array.shape)
+        data_str = json.dumps(array.flatten().tolist())
+        self.data_json = self._compress_bytes(data_str.encode())
+
+
+    def get_array(self):
+        """Retrieve the numpy array"""
+        import numpy as np
+
+        shape = tuple(int(dim) for dim in self.shape.split(",")) if self.shape else ()
+        flat_array = np.array([])
+        if self.data_json:
+            try:
+                # Try to decompress assuming it's compressed
+                data_str = self._decompress_bytes(self.data_json).decode()
+            except (zlib.error, Exception):
+                # If decompression fails, treat as uncompressed
+                data_str = self.data_json
+            flat_array = np.array(json.loads(data_str))
+        return flat_array.reshape(shape)
+
+    @property
+    def array(self):
+        """Property getter for array"""
+        return self.get_array()
+
+    @array.setter
+    def array(self, value):
+        """Property setter for array"""
+        self.set_array(value)
+
+    def make_table_entries(
+        self,
+        max_recursion_level=1,
+        drop_id=True,
+        current_level=0,
+        visited=None,
+        field_prefix="",
+    ):
+        return {"name": self.name}
+
+    def make_column_defs_instance(
+        self,
+        table_name=None,
+        max_recursion_level=1,
+        drop_id=True,
+        current_level=0,
+        visited=None,
+        field_prefix="",
+    ):
+        return [{"field": "name", "headerName": "Array"}]
+
+    @classmethod
+    def ui_schema(cls, **kwargs) -> dict:
+        return ui_hide_fields({}, ["shape", "field_name"])

simstack/models/artifact_models.py

@@ -0,0 +1,127 @@
+from typing import Optional, Dict, Any
+
+from odmantic import Field, Model
+
+from simstack.models.pickle_models import FunctionPickle
+from simstack.models.simstack_model import simstack_model
+
+
+# class ArtifactList(list):
+#     """A list that knows how to update its parent Artifact when modified."""
+#
+#     def __init__(self, parent_artifact, key, initial_values=None):
+#         super().__init__(initial_values or [])
+#         self.parent_artifact = parent_artifact
+#         self.key = key
+#
+#     def append(self, value):
+#         super().append(value)
+#         self.parent_artifact[self.key] = self  # Update the parent
+#
+#     def extend(self, values):
+#         super().extend(values)
+#         self.parent_artifact[self.key] = self  # Update the parent
+#
+#     # Add other list methods that modify the list
+#     def __setitem__(self, index, value):
+#         super().__setitem__(index, value)
+#         self.parent_artifact[self.key] = self
+#
+#     def __delitem__(self, index):
+#         super().__delitem__(index)
+#         self.parent_artifact[self.key] = self
+#
+#     def insert(self, index, value):
+#         super().insert(index, value)
+#         self.parent_artifact[self.key] = self
+#
+#     def pop(self, index=-1):
+#         value = super().pop(index)
+#         self.parent_artifact[self.key] = self
+#         return value
+#
+#     def remove(self, value):
+#         super().remove(value)
+#         self.parent_artifact[self.key] = self
+#
+#     def clear(self):
+#         super().clear()
+#         self.parent_artifact[self.key] = self
+#
+#     def sort(self, *args, **kwargs):
+#         super().sort(*args, **kwargs)
+#         self.parent_artifact[self.key] = self
+#
+#     def reverse(self):
+#         super().reverse()
+#         self.parent_artifact[self.key] = self
+
+
+@simstack_model
+class ArtifactModel(Model):
+    name: str
+    description: Optional[str] = None
+    data: Dict[str, Any] = Field(default_factory=dict)
+    path: Optional[str] = None
+
+    model_config = {
+        "collection": "artifacts"  # All subclasses will use the same collection
+    }
+
+
+# class IntArtifactModel(ArtifactModel):
+#     type: Literal["int"] = "int"  # Fixed value for this subclass
+#     value: int
+#
+# class FloatArtifactModel(ArtifactModel):
+#     type: Literal["float"] = "float"  # Fixed value for this subclass
+#     value: float
+#
+# class StringArtifactModel(ArtifactModel):
+#     type: Literal["string"] = "string"  # Fixed value for this subclass
+#     value: str
+#
+# class ImageArtifactModel(ArtifactModel):
+#     type: Literal["image"] = "image"  # Fixed value for this subclass
+#     value: str  # the image is stored as a base64 string or a file path
+#
+# # Forward reference for type hints
+# ArtifactType = Union[ArtifactModel, IntArtifactModel, FloatArtifactModel, StringArtifactModel, 'ArtifactModelList', ImageArtifactModel]
+#
+# # List subclass - can contain other items
+# class ArtifactModelList(ArtifactModel):
+#     type: Literal["list"] = "list"
+#     items: List[ArtifactType] = Field(default_factory=list)  # List of Item objects
+
+
+class ArtifactMapping(Model):
+    """
+    ArtifactsMapper is a mapping between the artifact and a node registry-path.
+    The workflow executor passes a path of the type
+
+    node1.node2.node4. ... .nodeN
+
+    where node is the function name of the node
+
+    Regex can maps this to the target path of the ArtifactsMapping, e.g. a path
+
+    *.parent1.node
+
+    it would map on all nodes with name node that have been directly called by a node with the name parent1.
+
+
+    """
+
+    name: str = Field(default="artifact", unique=True)
+    regex_pattern: str = Field(default="")
+    function_mapping: str = Field(default="")
+    function_code: str = Field(default="")
+    pickle_function: Optional[FunctionPickle] = None
+
+    def set_values(self, other: "ArtifactMapping") -> "ArtifactMapping":
+        self.name = other.name
+        self.regex_pattern = other.regex_pattern
+        self.function_mapping = other.function_mapping
+        self.function_code = other.function_code
+        self.pickle_function = other.pickle_function
+        return self

simstack/models/base_lists.py

@@ -0,0 +1,51 @@
+from typing import Iterator, List, TypeVar
+
+from odmantic import Field, Model, ObjectId
+
+from simstack.models.base_types import BooleanData, StringData
+from simstack.models.simstack_model import simstack_model
+from simstack.util.generic_list_mixin import GenericListMixin
+from simstack.util.object_list_mixin import ObjectListMixin
+
+T = TypeVar("T")
+
+@simstack_model
+class StringDataList(Model, ObjectListMixin[StringData]):
+    field_name: str = "string_data_list"
+    elements: List[ObjectId] = Field(default_factory=list, description="List of StringData ObjectIDs")
+
+    def __iter__(self) -> Iterator[StringData]:
+        return ObjectListMixin.__iter__(self)
+
+    def __init__(self, **data):
+        data, cache = self._normalize_elements_for_init(data)
+        Model.__init__(self, **data)
+        if cache is not None:
+            self._set_cache(cache)
+
+
+@simstack_model
+class BooleanDataList(Model, ObjectListMixin[BooleanData]):
+    field_name: str = "boolean_data_list"
+    elements: List[ObjectId] = Field(default_factory=list, description="List of BooleanData ObjectIDs")
+
+    def __iter__(self) -> Iterator[BooleanData]:
+        return ObjectListMixin.__iter__(self)
+
+    def __init__(self, **data):
+        data, cache = self._normalize_elements_for_init(data)
+        Model.__init__(self, **data)
+        if cache is not None:
+            self._set_cache(cache)
+
+
+@simstack_model
+class StringList(Model, GenericListMixin[str]):
+    field_name: str = "string_list"
+    elements: List[str] = Field(default_factory=list, description="List of strings")
+
+    def __iter__(self) -> Iterator[T]:
+        return iter(self.elements)
+
+    def __len__(self) -> int:
+        return len(self.elements)

simstack/models/base_types.py

@@ -0,0 +1,209 @@
+from odmantic import Model, Reference
+from pydantic import model_validator
+
+from simstack.models import simstack_model
+
+
+@simstack_model
+class IntData(Model):
+    field_name: str = "int"
+    value: int
+
+    @model_validator(mode="before")
+    @classmethod
+    def ensure_fieldname(cls, data):
+        """Ensure fieldname is set for existing documents"""
+        if isinstance(data, dict) and "field_name" not in data:
+            data["field_name"] = "int"
+        return data
+
+    def make_table_entries(
+        self,
+        max_recursion_level=1,
+        drop_id=True,
+        current_level=0,
+        visited=None,
+        field_prefix="",
+    ):
+        return {self.field_name: self.value}
+
+    def make_column_defs_instance(
+        self,
+        table_name=None,
+        max_recursion_level=1,
+        drop_id=True,
+        current_level=0,
+        visited=None,
+        field_prefix="",
+    ):
+        return [
+            {
+                "field": self.field_name,
+                "headerName": self.field_name,
+                "type": "numericColumn",
+            }
+        ]
+
+    @classmethod
+    def ui_schema(cls, **kwargs) -> dict:
+        return {
+            "ui:field": "IntDataField"
+        }  # This tells RJSF to use the custom component
+
+
+@simstack_model
+class FloatData(Model):
+    field_name: str = "float"
+    value: float
+
+    @model_validator(mode="before")
+    @classmethod
+    def ensure_fieldname(cls, data):
+        """Ensure fieldname is set for existing documents"""
+        if isinstance(data, dict) and "field_name" not in data:
+            data["field_name"] = "float"
+        return data
+
+    def make_table_entries(
+        self,
+        max_recursion_level=1,
+        drop_id=True,
+        current_level=0,
+        visited=None,
+        field_prefix="",
+    ):
+        return {self.field_name: self.value}
+
+    def make_column_defs_instance(
+        self,
+        table_name=None,
+        max_recursion_level=1,
+        drop_id=True,
+        current_level=0,
+        visited=None,
+        field_prefix="",
+    ):
+        return [
+            {
+                "field": self.field_name,
+                "headerName": self.field_name,
+                "type": "numericColumn",
+            }
+        ]
+
+    @classmethod
+    def ui_schema(cls, **kwargs) -> dict:
+        return {
+            "ui:field": "FloatDataField"
+        }  # This tells RJSF to use the custom component
+
+
+@simstack_model
+class StringData(Model):
+    field_name: str = "text"
+    value: str
+
+    @model_validator(mode="before")
+    @classmethod
+    def ensure_fieldname(cls, data):
+        """Ensure fieldname is set for existing documents"""
+        if isinstance(data, dict) and "field_name" not in data:
+            data["field_name"] = "text"
+        return data
+
+    def make_table_entries(
+        self,
+        max_recursion_level=1,
+        drop_id=True,
+        current_level=0,
+        visited=None,
+        field_prefix="",
+    ):
+        return {self.field_name: self.value}
+
+    def make_column_defs_instance(
+        self,
+        table_name=None,
+        max_recursion_level=1,
+        drop_id=True,
+        current_level=0,
+        visited=None,
+        field_prefix="",
+    ):
+        return [
+            {
+                "field": self.field_name,
+                "headerName": self.field_name,
+                "type": "textColumn",
+            }
+        ]
+
+    @classmethod
+    def ui_schema(cls, **kwargs) -> dict:
+        return {
+            "ui:field": "StringDataField"
+        }  # This tells RJSF to use the custom component
+
+
+@simstack_model
+class BooleanData(Model):
+    field_name: str = "boolean"
+    value: bool
+
+    @model_validator(mode="before")
+    @classmethod
+    def ensure_fieldname(cls, data):
+        """Ensure fieldname is set for existing documents"""
+        if isinstance(data, dict) and "field_name" not in data:
+            data["field_name"] = "boolean"
+        return data
+
+    def make_table_entries(
+        self,
+        max_recursion_level=1,
+        drop_id=True,
+        current_level=0,
+        visited=None,
+        field_prefix="",
+    ):
+        return {self.field_name: self.value}
+
+    def make_column_defs_instance(
+        self,
+        table_name=None,
+        max_recursion_level=1,
+        drop_id=True,
+        current_level=0,
+        visited=None,
+        field_prefix="",
+    ):
+        return [
+            {
+                "field": self.field_name,
+                "headerName": self.field_name,
+                "type": "booleanColumn",
+            }
+        ]
+
+    @classmethod
+    def ui_schema(cls, **kwargs) -> dict:
+        return {
+            "ui:field": "BooleanDataField"
+        }  # This tells RJSF to use the custom component
+
+
+@simstack_model
+class BinaryOperationInput(Model):
+    field_name: str = "binary_operation"
+    arg1: FloatData = Reference()
+    arg2: FloatData = Reference()
+
+    def __str__(self):
+        return f"BinaryOperationInput(arg1={self.arg1}, arg2={self.arg2})"
+
+
+@simstack_model
+class IteratorInput(Model):
+    start: int
+    stop: int
+    generator: str = "range"