py-modelrunner 0.16.0__tar.gz → 0.17.1__tar.gz

{py-modelrunner-0.16.0 → py_modelrunner-0.17.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: py-modelrunner
-Version: 0.16.0
+Version: 0.17.1
 Summary: Python classes for organizing (HPC) simulations
 Author-email: David Zwicker <david.zwicker@ds.mpg.de>
 License: MIT
@@ -226,7 +226,7 @@ def multiply(a=1, b=2):
     return a * b
 
 if __name__ == "__main__":
-    submit_job(__file__, parameters={'a': 2}, output="data.hdf5", method="foreground")
+    submit_job(__file__, parameters={"a": 2}, output="data.hdf5", method="foreground")
 ```
 Here, the `output` argument specifies a file to which the results are written, while
 `method` chooses how the script is submitted.
@@ -243,7 +243,7 @@ def multiply(a=1, b=2):
 
 if __name__ == "__main__":
     for a in range(5):
-        submit_job(__file__, parameters={'a': a}, output=f"data_{a}.hdf5", method="foreground")
+        submit_job(__file__, parameters={"a": a}, output=f"data_{a}.hdf5", method="foreground")
 ```
 
 Note that the safe-guard `if __name__ == "__main__"` is absolutely crucial to ensure that
@@ -261,7 +261,7 @@ def multiply(a=1, b=2):
     return a * b
 
 if __name__ == "__main__":
-    submit_jobs(__file__, parameters={'a': range(5)}, output_folder="data", method="foreground")
+    submit_jobs(__file__, parameters={"a": range(5)}, output_folder="data", method="foreground")
 ```
 
 Finally, the packages also offers a method to submit a model script to the cluster using
@@ -301,7 +301,7 @@ For instance, the data from the multiple jobs ran above can be collected using
 from modelrunner import ResultCollection
 
 results = ResultCollection.from_folder(".", pattern="data_*.hdf5")
-print(results.dataframe)
+print(results.to_dataframe())
 ```
 
 This example should print all results using a pandas Dataframe, where each row

{py-modelrunner-0.16.0 → py_modelrunner-0.17.1}/README.md

@@ -184,7 +184,7 @@ def multiply(a=1, b=2):
     return a * b
 
 if __name__ == "__main__":
-    submit_job(__file__, parameters={'a': 2}, output="data.hdf5", method="foreground")
+    submit_job(__file__, parameters={"a": 2}, output="data.hdf5", method="foreground")
 ```
 Here, the `output` argument specifies a file to which the results are written, while
 `method` chooses how the script is submitted.
@@ -201,7 +201,7 @@ def multiply(a=1, b=2):
 
 if __name__ == "__main__":
     for a in range(5):
-        submit_job(__file__, parameters={'a': a}, output=f"data_{a}.hdf5", method="foreground")
+        submit_job(__file__, parameters={"a": a}, output=f"data_{a}.hdf5", method="foreground")
 ```
 
 Note that the safe-guard `if __name__ == "__main__"` is absolutely crucial to ensure that
@@ -219,7 +219,7 @@ def multiply(a=1, b=2):
     return a * b
 
 if __name__ == "__main__":
-    submit_jobs(__file__, parameters={'a': range(5)}, output_folder="data", method="foreground")
+    submit_jobs(__file__, parameters={"a": range(5)}, output_folder="data", method="foreground")
 ```
 
 Finally, the packages also offers a method to submit a model script to the cluster using
@@ -259,7 +259,7 @@ For instance, the data from the multiple jobs ran above can be collected using
 from modelrunner import ResultCollection
 
 results = ResultCollection.from_folder(".", pattern="data_*.hdf5")
-print(results.dataframe)
+print(results.to_dataframe())
 ```
 
 This example should print all results using a pandas Dataframe, where each row

{py-modelrunner-0.16.0 → py_modelrunner-0.17.1}/docs/source/conf.py

@@ -240,7 +240,6 @@ napoleon_custom_sections = None
 intersphinx_mapping = {
     "h5py": ("https://docs.h5py.org/en/latest", None),
     "matplotlib": ("https://matplotlib.org/stable", None),
-    "modelrunner": ("https://py-modelrunner.readthedocs.io/en/latest", None),
     "napari": ("https://napari.org/", None),
     "numba": ("https://numba.pydata.org/numba-doc/latest/", None),
     "numpy": ("https://numpy.org/doc/stable", None),

{py-modelrunner-0.16.0 → py_modelrunner-0.17.1}/docs/source/user_manual.rst

@@ -43,11 +43,12 @@ The main requirements for the storage classes were
 - *Usability*: The user should not need to think about how data is stored in different files
 - *Flexibility*: We want a general interface to write data in multiple file formats (YAML, HDF, zarr, ...)
 - *Stability*: Future versions of the package should be able to read older files even when the internal definitions of file formats change
-- *Extensibility*: States should be subclasses to implement more complicated requirements (e.g., particular serialization)
+- *Modularity*: Different parts of the package (like :mod:`~modelrunner.storage`, :mod:`~modelrunner.parameters`, and :mod:`~modelrunner.run`) should be rather independent of each other, so they can be used in isolation
+- *Extensibility*: Models should be easy to subclass to implement more complicated requirements (e.g., additional parameters)
 - *Self-explainability*: The files should in principle contain all information to reconstruct the data, even if the `py-modelrunner` package is no longer available.
 - *Efficiency*: The files should only store necessary information.
 
 The last point results in particular constraints if we want to store temporal simulation results.
 In most cases, there are are some data that are kept fixed for the simulation (describing physical parameters) and others that evolve with time.
 We denote by `attributes` the parameters that are kept fixed and by `data` the data that varies over time.
-The state classes are already prepared to deal with such data, in conjuction with the :mod:`~modelrunner.storage.trajectory` module.
+The :mod:`~modelrunner.storage.trajectory` module deals with such data.

{py-modelrunner-0.16.0 → py_modelrunner-0.17.1}/examples/io_hdf.py

@@ -22,4 +22,4 @@ if __name__ == "__main__":
 
     # write result from file
     read = Result.from_file("test.hdf")
-    print(read.parameters, "–– start + [0..length-1] =", read.data)
+    print(read.parameters, "–– start + [0..length-1] =", read.result)

{py-modelrunner-0.16.0 → py_modelrunner-0.17.1}/examples/io_json.py

@@ -18,4 +18,4 @@ if __name__ == "__main__":
     result.to_file("test.json")
 
     read = Result.from_file("test.json")
-    print(read.parameters, "–– a * b =", read.data)
+    print(read.parameters, "–– a * b =", read.result)

{py-modelrunner-0.16.0 → py_modelrunner-0.17.1}/examples/io_yaml.py

@@ -18,4 +18,4 @@ if __name__ == "__main__":
     result.to_file("test.yaml")
 
     read = Result.from_file("test.yaml")
-    print(read.parameters, "–– a * b =", read.data)
+    print(read.parameters, "–– a * b =", read.result)

{py-modelrunner-0.16.0 → py_modelrunner-0.17.1}/examples/model_storage_output.py

@@ -24,4 +24,4 @@ with tempfile.NamedTemporaryFile(suffix=".yaml") as fp:
     # read the file and check whether all the data is there
     with open_storage(fp.name) as storage:
         print("Stored data:", storage["data"])
-        print("Model result:", storage["result"].data)
+        print("Model result:", storage["result"].result)

{py-modelrunner-0.16.0 → py_modelrunner-0.17.1}/examples/script_function.py

@@ -15,4 +15,4 @@ def multiply(a: float = 1, b: float = 2):
 
 if __name__ == "__main__":
     result = run_function_with_cmd_args(multiply)
-    print(result.data)
+    print(result.result)

py_modelrunner-0.17.1/modelrunner/.gitignore

@@ -0,0 +1 @@
+/_version.py

{py-modelrunner-0.16.0 → py_modelrunner-0.17.1}/modelrunner/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.16.0'
-__version_tuple__ = version_tuple = (0, 16, 0)
+__version__ = version = '0.17.1'
+__version_tuple__ = version_tuple = (0, 17, 1)

{py-modelrunner-0.16.0 → py_modelrunner-0.17.1}/modelrunner/model.py

@@ -64,7 +64,11 @@ class ModelBase(Parameterized, metaclass=ABCMeta):
                 :meth:`~Parameterized.get_parameters` or displayed by calling
                 :meth:`~Parameterized.show_parameters`.
             output (str):
-                Path where the output file will be written.
+                Path where the output file will be written. The output will be written
+                using :mod:`~modelrunner.storage` and might contain two groups: `result`
+                to which the final result of the model is written, and `data`, which
+                can contain extra information that is written using
+                :attr:`~ModelBase.storage`.
             mode (str or :class:`~modelrunner.storage.access_modes.ModeType`):
                 The file mode with which the storage is accessed, which determines the
                 allowed operations. Common options are "read", "full", "append", and
@@ -87,7 +91,9 @@ class ModelBase(Parameterized, metaclass=ABCMeta):
             if self.output is None:
                 raise RuntimeError("Output file needs to be specified")
             self._storage = open_storage(self.output, mode=self.mode)
-        return self._storage
+            return self._storage.create_group("data")
+        else:
+            return self._storage.open_group("data")
 
     def close(self) -> None:
         """close any opened storages"""
@@ -124,21 +130,21 @@ class ModelBase(Parameterized, metaclass=ABCMeta):
             result:
                 The result data. If omitted, the model is run to obtain results
         """
+        from .results import Result  # @Reimport
+
         if self.output is None:
             raise RuntimeError("Output file needs to be specified")
 
         if result is None:
             result = self.get_result()
-        else:
-            from .results import Result  # @Reimport
-
-            assert isinstance(result, Result)
+        elif not isinstance(result, Result):
+            raise TypeError(f"result has type {result.__class__} instead of `Result`")
 
         if self._storage is not None:
             # reuse the opened storage
-            result.to_file(self.storage)
+            result.to_file(self._storage, loc="result")
         else:
-            result.to_file(self.output, mode=self.mode)
+            result.to_file(self.output, loc="result", mode=self.mode)
 
     @classmethod
     def _prepare_argparser(cls, name: str | None = None) -> argparse.ArgumentParser:

{py-modelrunner-0.16.0 → py_modelrunner-0.17.1}/modelrunner/results.py

@@ -47,7 +47,12 @@ class Result:
     """int: number indicating the version of the file format"""
 
     def __init__(
-        self, model: ModelBase, result: Any, info: dict[str, Any] | None = None
+        self,
+        model: ModelBase,
+        result: Any,
+        *,
+        storage: StorageID = None,
+        info: dict[str, Any] | None = None,
     ):
         """
         Args:
@@ -55,18 +60,23 @@ class Result:
                 The model from which the result was obtained
             result:
                 The actual result
+            storage:
+                A storage containing additional data from the model run
             info (dict):
                 Additional information for this result
         """
         if not isinstance(model, ModelBase):
             raise TypeError("The model should be of type `ModelBase`")
-        self.model = model
         self.result = result
+        self.model = model
+        self.storage = storage
         self.info: Attrs = {} if info is None else info
 
     @property
     def data(self):
         """direct access to the underlying state data"""
+        # deprecated on 2024-04-13
+        warnings.warn("`.data` attribute was renamed to `.result`", DeprecationWarning)
         return self.result
 
     @classmethod
@@ -74,7 +84,9 @@ class Result:
         cls,
         model_data: dict[str, Any],
         result,
+        *,
         model: ModelBase | None = None,
+        storage: StorageID = None,
         info: dict[str, Any] | None = None,
     ) -> Result:
         """create result from data
@@ -86,6 +98,8 @@ class Result:
                 The actual result data
             model (:class:`ModelBase`):
                 The model from which the result was obtained
+            storage:
+                A storage containing additional data from the model run
             info (dict):
                 Additional information for this result
 
@@ -103,7 +117,7 @@ class Result:
         model.name = model_data.get("name")
         model.description = model_data.get("description")
 
-        return cls(model, result, info)
+        return cls(model, result, storage=storage, info=info)
 
     @property
     def parameters(self) -> dict[str, Any]:
@@ -119,13 +133,20 @@ class Result:
     ):
         """load object from a file
 
+        This function loads the results from a hierachical storage. It also attempts to
+        read information about the model that was used to create this result and
+        additional data that might have been stored in a
+        :attr:`~modelrunner.results.Result.storage` while the model was running.
+
         Args:
             store (str or :class:`zarr.Store`):
                 Path or instance describing the storage, which is either a file path or
                 a :class:`zarr.Storage`.
-            key (str):
-                Name of the node in which the data was stored. This applies to some
-                hierarchical storage formats.
+            loc:
+                The location where the result is stored in the storage. This should
+                rarely be modified.
+            model (:class:`~modelrunner.model.ModelBase`):
+                The model which lead to this result
         """
         if isinstance(storage, (str, Path)):
             # check whether the file was written with an old format version
@@ -141,10 +162,15 @@ class Result:
             format_version = attrs.pop("format_version", None)
             if format_version == cls._format_version:
                 # current version of storing results
+                if "data" in storage_obj:
+                    data_storage = open_storage(storage, loc="data", mode="read")
+                else:
+                    data_storage = None
                 return cls.from_data(
                     model_data=attrs.get("model", {}),
                     result=storage_obj.read_item(loc, use_class=False),
                     model=model,
+                    storage=data_storage,
                     info=attrs.pop("info", {}),  # load additional info,
                 )
 
@@ -154,7 +180,10 @@ class Result:
     def to_file(
         self, storage: StorageID, loc: str = "result", *, mode: ModeType = "insert"
     ) -> None:
-        """write this object to a file
+        """write the results to a file
+
+        Note that this does only write the actual `results` but omits additional data
+        that might have been stored in a storage that is associated with the results.
 
         Args:
             storage (:class:`StorageBase` or :class:`StorageGroup`):

{py-modelrunner-0.16.0 → py_modelrunner-0.17.1}/modelrunner/run/__main__.py

@@ -27,6 +27,12 @@ def main():
     )
 
     parser.add_argument("-o", "--output", help="Path to output file", metavar="PATH")
+    parser.add_argument(
+        "-l",
+        "--log_folder",
+        help="Path for log files (skip for no logging)",
+        metavar="PATH",
+    )
     parser.add_argument(
         "--overwrite",
         nargs="?",
@@ -52,6 +58,7 @@ def main():
         output=args.output,
         name=args.name,
         parameters=args.parameters,
+        log_folder=args.log_folder,
         method=args.method,
         template=args.template,
         overwrite_strategy=args.overwrite,

{py-modelrunner-0.16.0 → py_modelrunner-0.17.1}/modelrunner/run/job.py

@@ -11,7 +11,7 @@ import itertools
 import json
 import logging
 import os
-import pipes
+import shlex
 import subprocess as sp
 import sys
 import warnings
@@ -26,7 +26,7 @@ from ..parameters import Parameter
 
 def escape_string(obj) -> str:
     """escape a string for the command line"""
-    return pipes.quote(str(obj))
+    return shlex.quote(str(obj))
 
 
 def ensure_directory_exists(folder):
@@ -130,7 +130,7 @@ def submit_job(
     parameters: str | dict[str, Any] | None = None,
     config: str | dict[str, Any] | None = None,
     *,
-    log_folder: str | Path = "logs",
+    log_folder: str | Path | None = None,
     method: str = "auto",
     use_modelrunner: bool = True,
     template: str | Path | None = None,
@@ -153,7 +153,10 @@ def submit_job(
             Configuration for the job, which determines how the job is run. Can be either
             a python dictionary or a string containing a JSON-encoded dictionary.
         log_folder (str of :class:`~pathlib.Path`):
-            Path to the logging folder
+            Path to the logging folder. If omitted, the default of the template is used,
+            which typically sends data to stdout for local scripts (which is thus
+            captured and returned by this function) or writes log files to the current
+            working directory for remote jobs.
         method (str):
             Specifies the submission method. Currently `background`, `foreground`,
             'srun', and `qsub` are supported. The special value `auto` reads the method
@@ -169,7 +172,9 @@ def submit_job(
             `error`, `warn_skip`, `silent_skip`, `overwrite`, and `silent_overwrite`.
 
     Returns:
-        tuple: The result `(stdout, stderr)` of the submission call
+        tuple: The result `(stdout, stderr)` of the submission call. These two strings
+            can contain the output from the actual scripts that is run when `log_folder`
+            is `None`.
     """
     from jinja2 import Template
 
@@ -191,7 +196,7 @@ def submit_job(
 
     # load the correct template
     if template is None:
-        template_path = Path(__file__).parent / "templates" / (method + ".template")
+        template_path = Path(__file__).parent / "templates" / (method + ".jinja")
     else:
         template_path = Path(template)
     logger.info("Load template `%s`", template_path)
@@ -199,15 +204,16 @@ def submit_job(
         script_template = fp.read()
 
     # prepare submission script
-    ensure_directory_exists(log_folder)
     script_args: dict[str, Any] = {
         "PACKAGE_PATH": Path(__file__).parents[2],
-        "LOG_FOLDER": log_folder,
         "JOB_NAME": name,
         "MODEL_FILE": escape_string(script),
         "USE_MODELRUNNER": use_modelrunner,
         "CONFIG": configuration,
     }
+    if log_folder is not None:
+        ensure_directory_exists(log_folder)
+        script_args["LOG_FOLDER"] = log_folder
 
     # add the parameters to the job arguments
     job_args = []
@@ -244,9 +250,9 @@ def submit_job(
 
         # check whether output points to a directory or whether this should be a file
         if output.is_dir():
-            script_args["OUTPUT_FOLDER"] = pipes.quote(str(output))
+            script_args["OUTPUT_FOLDER"] = shlex.quote(str(output))
         else:
-            script_args["OUTPUT_FOLDER"] = pipes.quote(str(output.parent))
+            script_args["OUTPUT_FOLDER"] = shlex.quote(str(output.parent))
             job_args.append(f"--output {escape_string(output)}")
 
     else:

py_modelrunner-0.17.1/modelrunner/run/templates/background.jinja

@@ -0,0 +1,27 @@
+#!/bin/bash -l
+
+export PYTHONPATH={{ PACKAGE_PATH }}:$PYTHONPATH
+
+{% if CONFIG.num_threads is number %}
+# set the number of threads to use
+export MKL_NUM_THREADS={{ CONFIG.num_threads }}
+export NUMBA_NUM_THREADS={{ CONFIG.num_threads }}
+export NUMEXPR_NUM_THREADS={{ CONFIG.num_threads }}
+export OMP_NUM_THREADS={{ CONFIG.num_threads }}
+export OPENBLAS_NUM_THREADS={{ CONFIG.num_threads }}
+{% endif %}
+
+{% if OUTPUT_FOLDER is defined and OUTPUT_FOLDER %}
+mkdir -p {{ OUTPUT_FOLDER }}
+{% endif %}
+
+# Run the program in the background
+nohup {{ CONFIG.python_bin }} \
+    {% if USE_MODELRUNNER is defined and USE_MODELRUNNER %} -m modelrunner {% endif %} \
+    {{ MODEL_FILE }} \
+    {{ JOB_ARGS }} \
+{% if LOG_FOLDER is defined -%}
+    > {{ LOG_FOLDER }}/{{ JOB_NAME }}.out.txt \
+    2> {{ LOG_FOLDER }}/{{ JOB_NAME }}.err.txt \
+{% endif -%}
+    &

py_modelrunner-0.17.1/modelrunner/run/templates/foreground.jinja

@@ -0,0 +1,26 @@
+#!/bin/bash -l
+
+export PYTHONPATH={{ PACKAGE_PATH }}:$PYTHONPATH
+
+{% if CONFIG.num_threads is number %}
+# set the number of threads to use
+export MKL_NUM_THREADS={{ CONFIG.num_threads }}
+export NUMBA_NUM_THREADS={{ CONFIG.num_threads }}
+export NUMEXPR_NUM_THREADS={{ CONFIG.num_threads }}
+export OMP_NUM_THREADS={{ CONFIG.num_threads }}
+export OPENBLAS_NUM_THREADS={{ CONFIG.num_threads }}
+{% endif %}
+
+{% if OUTPUT_FOLDER is defined and OUTPUT_FOLDER %}
+mkdir -p {{ OUTPUT_FOLDER }}
+{% endif %}
+
+# Run the program
+{{ CONFIG.python_bin }} \
+    {% if USE_MODELRUNNER is defined and USE_MODELRUNNER %} -m modelrunner {% endif %} \
+    {{ MODEL_FILE }} \
+    {{ JOB_ARGS }} \
+{% if LOG_FOLDER is defined -%}
+    > {{ LOG_FOLDER }}/{{ JOB_NAME }}.out.txt \
+    2> {{ LOG_FOLDER }}/{{ JOB_NAME }}.err.txt \
+{% endif -%}

py-modelrunner-0.16.0/modelrunner/run/templates/qsub.template → py_modelrunner-0.17.1/modelrunner/run/templates/qsub.jinja

@@ -18,13 +18,13 @@
 hostname
 echo $JOB_ID
 
-{% if CONFIG['num_threads'] is number %}
+{% if CONFIG.num_threads is number %}
 # set the number of threads to use
-export MKL_NUM_THREADS={{ CONFIG['num_threads'] }}
-export NUMBA_NUM_THREADS={{ CONFIG['num_threads'] }}
-export NUMEXPR_NUM_THREADS={{ CONFIG['num_threads'] }}
-export OMP_NUM_THREADS={{ CONFIG['num_threads'] }}
-export OPENBLAS_NUM_THREADS={{ CONFIG['num_threads'] }}
+export MKL_NUM_THREADS={{ CONFIG.num_threads }}
+export NUMBA_NUM_THREADS={{ CONFIG.num_threads }}
+export NUMEXPR_NUM_THREADS={{ CONFIG.num_threads }}
+export OMP_NUM_THREADS={{ CONFIG.num_threads }}
+export OPENBLAS_NUM_THREADS={{ CONFIG.num_threads }}
 {% endif %}
 
 {% if OUTPUT_FOLDER is defined and OUTPUT_FOLDER %}
@@ -33,7 +33,7 @@ mkdir -p {{ OUTPUT_FOLDER }}
 
 # Run the program
 {% if USE_MODELRUNNER is defined and USE_MODELRUNNER %}
-{{ CONFIG['python_bin'] }} -m modelrunner {{ MODEL_FILE }} {{ JOB_ARGS }}
+{{ CONFIG.python_bin }} -m modelrunner {{ MODEL_FILE }} {{ JOB_ARGS }}
 {% else %}
-{{ CONFIG['python_bin'] }} {{ MODEL_FILE }} {{ JOB_ARGS }}
+{{ CONFIG.python_bin }} {{ MODEL_FILE }} {{ JOB_ARGS }}
 {% endif %}

py-modelrunner-0.16.0/modelrunner/run/templates/srun.template → py_modelrunner-0.17.1/modelrunner/run/templates/srun.jinja

@@ -16,13 +16,13 @@
 hostname
 echo $JOB_ID
 
-{% if CONFIG['num_threads'] is number %}
+{% if CONFIG.num_threads is number %}
 # set the number of threads to use
-export MKL_NUM_THREADS={{ CONFIG['num_threads'] }}
-export NUMBA_NUM_THREADS={{ CONFIG['num_threads'] }}
-export NUMEXPR_NUM_THREADS={{ CONFIG['num_threads'] }}
-export OMP_NUM_THREADS={{ CONFIG['num_threads'] }}
-export OPENBLAS_NUM_THREADS={{ CONFIG['num_threads'] }}
+export MKL_NUM_THREADS={{ CONFIG.num_threads }}
+export NUMBA_NUM_THREADS={{ CONFIG.num_threads }}
+export NUMEXPR_NUM_THREADS={{ CONFIG.num_threads }}
+export OMP_NUM_THREADS={{ CONFIG.num_threads }}
+export OPENBLAS_NUM_THREADS={{ CONFIG.num_threads }}
 {% endif %}
 
 {% if OUTPUT_FOLDER is defined and OUTPUT_FOLDER %}
@@ -31,7 +31,7 @@ mkdir -p {{ OUTPUT_FOLDER }}
 
 # Run the program
 {% if USE_MODELRUNNER is defined and USE_MODELRUNNER %}
-{{ CONFIG['python_bin'] }} -m modelrunner {{ MODEL_FILE }} {{ JOB_ARGS }}
+{{ CONFIG.python_bin }} -m modelrunner {{ MODEL_FILE }} {{ JOB_ARGS }}
 {% else %}
-{{ CONFIG['python_bin'] }} {{ MODEL_FILE }} {{ JOB_ARGS }}
+{{ CONFIG.python_bin }} {{ MODEL_FILE }} {{ JOB_ARGS }}
 {% endif %}

{py-modelrunner-0.16.0 → py_modelrunner-0.17.1}/modelrunner/storage/backend/memory.py

@@ -56,17 +56,20 @@ class MemoryStorage(StorageBase):
             (group, str):
                 A tuple consisting of the parent group and the name of the current item
         """
-        value = self._data
+        if check_write and not self.is_group(loc[:-1]):
+            raise TypeError(f"Location `/{'/'.join(loc[:-1])}` is not a group")
+
+        parent = self._data
         for part in loc[:-1]:
             try:
-                value = value[part]
+                parent = parent[part]
             except KeyError:
-                if isinstance(value, dict):
-                    value[part] = {}
-                    value = value[part]
+                if isinstance(parent, dict):
+                    parent[part] = {}
+                    parent = parent[part]
                 else:
                     raise TypeError(f"Cannot add item to `/{'/'.join(loc)}`")
-        if not isinstance(value, dict):
+        if not isinstance(parent, dict):
             raise TypeError(f"Cannot add item to `/{'/'.join(loc)}`")
 
         try:
@@ -74,14 +77,17 @@ class MemoryStorage(StorageBase):
         except IndexError:
             raise KeyError(f"Location `/{'/'.join(loc)}` has no parent")
 
-        if check_write and not self.mode.overwrite and name in value:
+        if check_write and not self.mode.overwrite and name in parent:
             raise AccessError(f"Overwriting `/{'/'.join(loc)}` disabled")
-        return value, name
+        return parent, name
 
     def __getitem__(self, loc: Sequence[str]) -> Any:
         if loc:
             parent, name = self._get_parent(loc)
-            return parent[name]
+            try:
+                return parent[name]
+            except KeyError as e:
+                raise KeyError(f"Item `{name}` not in group {parent.keys()}") from e
         else:
             return self._data