wrfrun 0.2.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

wrfrun/core/_record.py

@@ -0,0 +1,178 @@
+"""
+wrfrun.core._record
+###################
+
+.. autosummary::
+    :toctree: generated/
+
+    ExecutableRecorder
+
+ExecutableRecorder
+******************
+
+This class provides methods to record simulations.
+It will save configs and resources of ``Executable``, and input data optionally.
+A file ends with ``.replay`` will be generated after finishing recording,
+and users can reproduce the simulation with the ``.replay`` file.
+"""
+
+from json import dumps
+from os import makedirs, remove
+from os.path import basename, dirname, exists, isdir
+from shutil import copyfile, make_archive, move
+
+import numpy as np
+
+from ..log import check_path, logger
+from ._config import WRFRunConfig
+from .type import ExecutableConfig
+
+
+def _json_default(obj):
+    """
+    Used for json.dumps.
+
+    :param obj:
+    :type obj:
+    :return:
+    :rtype:
+    """
+    if isinstance(obj, np.integer):
+        return int(obj)
+    elif isinstance(obj, np.floating):
+        return float(obj)
+    else:
+        raise TypeError(f"Object of type {type(obj)} is not JSON serializable.")
+
+
+class ExecutableRecorder:
+    """
+    This class provides methods to record simulations.
+    """
+
+    def __init__(self, wrfrun_config: WRFRunConfig, save_path="./wrfrun.replay", include_data=False):
+        """
+        :param wrfrun_config: `WRFRunConfig` instance.
+        :type wrfrun_config: WRFRunConfig
+        :param save_path: Save path of the replay file, defaults to "./wrfrun.replay"
+        :type save_path: str, optional
+        :param include_data: If includes data files, defaults to False
+        :type include_data: bool, optional
+        """
+        self._wrfrun_config = wrfrun_config
+
+        self.save_path = save_path
+        self.include_data = include_data
+
+        self.work_path = self._wrfrun_config.parse_resource_uri(self._wrfrun_config.WRFRUN_WORKSPACE_REPLAY)
+        self.content_path = f"{self.work_path}/config_and_data"
+
+        self._recorded_config = []
+        self._name_count = {}
+
+    def record(self, exported_config: ExecutableConfig):
+        """
+        Record exported config for replay.
+
+        :param exported_config: Executable config.
+        :type exported_config: ExecutableConfig
+        """
+        if not self.include_data:
+            self._recorded_config.append(exported_config)
+            return
+
+        check_path(self.content_path)
+
+        # process exported config so we can also include data.
+        # create directory to place data
+        name = exported_config["name"]
+        if name in self._name_count:
+            self._name_count[name] += 1
+            index = self._name_count[name]
+        else:
+            self._name_count[name] = 1
+            index = 1
+
+        data_save_uri = f"{self._wrfrun_config.WRFRUN_WORKSPACE_REPLAY}/{name}/{index}"
+        data_save_path = f"{self.content_path}/{name}/{index}"
+        makedirs(data_save_path)
+
+        input_file_config = exported_config["input_file_config"]
+
+        for _config_index, _config in enumerate(input_file_config):
+            if not _config["is_data"]:
+                continue
+
+            if _config["is_output"]:
+                continue
+
+            file_path = _config["file_path"]
+            file_path = self._wrfrun_config.parse_resource_uri(file_path)
+            filename = basename(file_path)
+            copyfile(file_path, f"{data_save_path}/{filename}")
+
+            _config["file_path"] = f"{data_save_uri}/{filename}"
+            input_file_config[_config_index] = _config
+
+        exported_config["input_file_config"] = input_file_config
+        self._recorded_config.append(exported_config)
+
+    def clear_records(self):
+        """
+        Clean recorded configs.
+        """
+        self._recorded_config = []
+
+    def set_recorder(self, save_path: str | None, include_data: bool | None):
+        """
+        Change recorder settings.
+
+        :param save_path: Save path of the exported config file.
+        :type save_path: str | None
+        :param include_data: If includes input data.
+        :type include_data: bool | None
+        """
+        if save_path is not None:
+            logger.debug(f"Change save path to: {save_path}")
+            self.save_path = save_path
+
+        if include_data is not None:
+            self.include_data = include_data
+
+    def export_replay_file(self):
+        """
+        Save replay file to the save path.
+        """
+        if len(self._recorded_config) == 0:
+            logger.warning("No replay config has been recorded.")
+            return
+
+        logger.info("Exporting replay config... It may take a few minutes if you include data.")
+
+        check_path(self.content_path)
+
+        with open(f"{self.content_path}/config.json", "w") as f:
+            f.write(dumps(self._recorded_config, indent=4, default=_json_default))
+
+        if exists(self.save_path):
+            if isdir(self.save_path):
+                self.save_path = f"{self.save_path}/wrfrun.replay"
+            else:
+                if not self.save_path.endswith(".replay"):
+                    self.save_path = f"{self.save_path}.replay"
+
+                if exists(self.save_path):
+                    logger.warning(f"Found existed replay file with the same name '{basename(self.save_path)}', overwrite it")
+                    remove(self.save_path)
+
+        if not exists(dirname(self.save_path)):
+            makedirs(dirname(self.save_path))
+
+        temp_file = f"{self.work_path}/config_and_data"
+        make_archive(temp_file, "zip", self.content_path)
+        move(f"{temp_file}.zip", self.save_path)
+
+        logger.info(f"Replay config exported to {self.save_path}")
+
+
+__all__ = ["ExecutableRecorder"]

wrfrun/core/_resource.py

@@ -0,0 +1,172 @@
+"""
+wrfrun.core._resource
+#####################
+
+.. autosummary::
+    :toctree: generated/
+
+    ResourceMixIn
+
+ResourceMixIn
+*************
+
+This class provides methods to manage resources in wrfrun.
+
+URIs in wrfrun
+==============
+
+``wrfrun`` uses URI (Uniform Resource Identifier) to represents real resource files,
+to make sure the same code works on different machines.
+
+There are some additional restrictions about URIs used in wrfrun:
+
+* Must start with ``:WRFRUN_``
+* Must end with ``:``
+
+Register URIs
+=============
+
+There are two ways to register URI:
+
+**Register directly**
+
+The most convenient way to register resource uri is using :meth:`register_resource_uri <ResourceMixIn.register_resource_uri>`.
+
+.. code-block:: Python
+    :caption: main.py
+
+    resource = ResourceMixIn()
+    resource_uri = ":WRFRUN_TEST_URI:"
+
+    # remember to check if it has been registered.
+    if not resource.check_resource_uri(resource_uri):
+        resource.register_resource_uri(resource_uri)
+
+**Use register function**
+
+Please see documentation about :class:`WRFRUN <wrfrun.core.core._WRFRUNProxy>`.
+
+Parse URIs
+==========
+
+You can get the real file path using :meth:`parse_resource_uri <ResourceMixIn.parse_resource_uri>`,
+it will parse all the URIs in the string.
+"""
+
+from ..log import logger
+from .error import ResourceURIError
+
+
+class ResourceMixIn:
+    """
+    This class provides methods to manage resources in wrfrun.
+
+    These resources may include various configuration files from NWP as well as those provided by wrfrun itself.
+    Since their actual file paths may vary depending on the installation environment,
+    wrfrun maps them using URIs to ensure consistent access regardless of the environment.
+
+    The URI always starts with the prefix string ``:WRFRUN_`` and ends with ``:``.
+
+    To register URIs, user can use :meth:`ResourceMixIn.register_resource_uri`.
+
+    To convert any possible URIs in a string, user can use :meth:`ResourceMixIn.parse_resource_uri`.
+    """
+
+    def __init__(self, *args, **kwargs):
+        self._resource_namespace_db = {}
+
+        super().__init__(*args, **kwargs)
+
+    def check_resource_uri(self, unique_uri: str) -> bool:
+        """
+        Check if the URI has been registered.
+
+        ``wrfrun`` uses unique URIs to represent resource files.
+        If you want to register a URI, you need to check if it's available.
+
+        :param unique_uri: Unique URI represents the resource.
+        :type unique_uri: str
+        :return: True or False.
+        :rtype: bool
+        """
+        if unique_uri in self._resource_namespace_db:
+            return True
+
+        else:
+            return False
+
+    def register_resource_uri(self, unique_uri: str, res_space_path: str):
+        """
+        Register a resource path with a URI.
+        The URI should start with ``:WRFRUN_`` ,end with ``:`` and hasn't been registered yet,
+        otherwise an exception :class:`ResourceURIError <wrfrun.core.error.ResourceURIError>` will be raised.
+
+        :param unique_uri: Unique URI represents the resource.
+                           It must start with ``:WRFRUN_`` and end with ``:``. For example, ``":WRFRUN_WORK_PATH:"``.
+        :type unique_uri: str
+        :param res_space_path: REAL absolute path of your resource path. For example, "$HOME/.config/wrfrun/res".
+        :type res_space_path: str
+        """
+        if not (unique_uri.startswith(":WRFRUN_") and unique_uri.endswith(":")):
+            logger.error(f"Can't register resource URI: '{unique_uri}'. It should start with ':WRFRUN_' and end with ':'.")
+            raise ResourceURIError(
+                f"Can't register resource URI: '{unique_uri}'. It should start with ':WRFRUN_' and end with ':'."
+            )
+
+        if unique_uri in self._resource_namespace_db:
+            logger.error(f"Resource URI '{unique_uri}' exists.")
+            raise ResourceURIError(f"Resource URI '{unique_uri}' exists.")
+
+        logger.debug(f"Register URI '{unique_uri}' to '{res_space_path}'")
+        self._resource_namespace_db[unique_uri] = res_space_path
+
+    def unregister_resource_uri(self, unique_uri: str):
+        """
+        Unregister a resource URI.
+
+        :param unique_uri: Unique URI represents the resource.
+        :type unique_uri: str
+        """
+        if unique_uri in self._resource_namespace_db:
+            self._resource_namespace_db.pop(unique_uri)
+
+    def parse_resource_uri(self, resource_path: str) -> str:
+        """
+        Return the converted string by parsing the URI string in it.
+        Normal path will be returned with no change.
+
+        If the URI hasn't been registered, an exception :class:`ResourceURIError` will be raised.
+
+        For example, you can get the real path of ``wrfrun`` workspace with this method:
+
+        >>> from wrfrun.core import WRFRUN
+        >>> workspace_path = f"{WRFRUN.config.WRFRUN_WORKSPACE_ROOT}/WPS"    # ":WRFRUN_WORKSPACE_PATH:/WPS"
+        >>> # real_path should be a valid path like: "/home/syize/.config/wrfrun/workspace/WPS"
+        >>> real_path = WRFRUN.config.parse_resource_uri(workspace_path)
+
+        :param resource_path: Resource path string which may contain URI string.
+        :type resource_path: str
+        :return: Real resource path.
+        :rtype: str
+        """
+        if not resource_path.startswith(":WRFRUN_"):
+            return resource_path
+
+        res_namespace_string = resource_path.split(":")[1]
+        res_namespace_string = f":{res_namespace_string}:"
+
+        if res_namespace_string in self._resource_namespace_db:
+            resource_path = resource_path.replace(res_namespace_string, self._resource_namespace_db[res_namespace_string])
+
+            if not resource_path.startswith(":WRFRUN_"):
+                return resource_path
+
+            else:
+                return self.parse_resource_uri(resource_path)
+
+        else:
+            logger.error(f"Unknown resource URI: '{res_namespace_string}'")
+            raise ResourceURIError(f"Unknown resource URI: '{res_namespace_string}'")
+
+
+__all__ = ["ResourceMixIn"]