pastastore 1.4.0__py3-none-any.whl → 1.6.0__py3-none-any.whl

pastastore/connectors.py

@@ -1,8 +1,9 @@
+"""Module containing classes for connecting to different data stores."""
+
 import json
 import os
 import warnings
 from copy import deepcopy
-from importlib import import_module
 from typing import Dict, Optional, Union
 
 import pandas as pd
@@ -15,190 +16,9 @@ FrameorSeriesUnion = Union[pd.DataFrame, pd.Series]
 warnings.showwarning = _custom_warning
 
 
-class ArcticConnector(BaseConnector, ConnectorUtil):  # pragma: no cover
-    conn_type = "arctic"
-
-    def __init__(self, name: str, connstr: str):
-        """Create an ArcticConnector object that connects to a running MongoDB
-        database via Arctic.
-
-        Parameters
-        ----------
-        name : str
-            name of the database
-        connstr : str
-            connection string (e.g. 'mongodb://localhost:27017/')
-        """
-        try:
-            import arctic
-        except ModuleNotFoundError as e:
-            print(
-                "Please install arctic (also requires "
-                "a MongoDB instance running somewhere, e.g. "
-                "MongoDB Community: \n"
-                "https://docs.mongodb.com/manual/administration"
-                "/install-community/)!"
-            )
-            raise e
-        self.connstr = connstr
-        self.name = name
-
-        self.libs: dict = {}
-        self.arc = arctic.Arctic(connstr)
-        self._initialize()
-        self.models = ModelAccessor(self)
-        # for older versions of PastaStore, if oseries_models library is empty
-        # populate oseries - models database
-        self._update_all_oseries_model_links()
-
-    def _initialize(self) -> None:
-        """Internal method to initalize the libraries."""
-
-        for libname in self._default_library_names:
-            if self._library_name(libname) not in self.arc.list_libraries():
-                self.arc.initialize_library(self._library_name(libname))
-            else:
-                print(
-                    f"ArcticConnector: library "
-                    f"'{self._library_name(libname)}'"
-                    " already exists. Linking to existing library."
-                )
-            self.libs[libname] = self._get_library(libname)
-
-    def _library_name(self, libname: str) -> str:
-        """Internal method to get full library name according to Arctic."""
-        return ".".join([self.name, libname])
-
-    def _get_library(self, libname: str):
-        """Get Arctic library handle.
-
-        Parameters
-        ----------
-        libname : str
-            name of the library
-
-        Returns
-        -------
-        lib : arctic.Library handle
-            handle to the library
-        """
-        # get library handle
-        lib = self.arc.get_library(self._library_name(libname))
-        return lib
-
-    def _add_item(
-        self,
-        libname: str,
-        item: Union[FrameorSeriesUnion, Dict],
-        name: str,
-        metadata: Optional[Dict] = None,
-        **_,
-    ) -> None:
-        """Internal method to add item to library (time series or model).
-
-        Parameters
-        ----------
-        libname : str
-            name of the library
-        item : Union[FrameorSeriesUnion, Dict]
-            item to add, either time series or pastas.Model as dictionary
-        name : str
-            name of the item
-        metadata : Optional[Dict], optional
-            dictionary containing metadata, by default None
-        """
-        lib = self._get_library(libname)
-        lib.write(name, item, metadata=metadata)
-
-    def _get_item(self, libname: str, name: str) -> Union[FrameorSeriesUnion, Dict]:
-        """Internal method to retrieve item from library.
-
-        Parameters
-        ----------
-        libname : str
-            name of the library
-        name : str
-            name of the item
-
-        Returns
-        -------
-        item : Union[FrameorSeriesUnion, Dict]
-            time series or model dictionary
-        """
-        lib = self._get_library(libname)
-        return lib.read(name).data
-
-    def _del_item(self, libname: str, name: str) -> None:
-        """Internal method to delete items (series or models).
-
-        Parameters
-        ----------
-        libname : str
-            name of library to delete item from
-        name : str
-            name of item to delete
-        """
-        lib = self._get_library(libname)
-        lib.delete(name)
-
-    def _get_metadata(self, libname: str, name: str) -> dict:
-        """Internal method to retrieve metadata for an item.
-
-        Parameters
-        ----------
-        libname : str
-            name of the library
-        name : str
-            name of the item
-
-        Returns
-        -------
-        dict
-            dictionary containing metadata
-        """
-        lib = self._get_library(libname)
-        return lib.read_metadata(name).metadata
-
-    @property
-    def oseries_names(self):
-        """List of oseries names.
-
-        Returns
-        -------
-        list
-            list of oseries in library
-        """
-        return self._get_library("oseries").list_symbols()
-
-    @property
-    def stresses_names(self):
-        """List of stresses names.
-
-        Returns
-        -------
-        list
-            list of stresses in library
-        """
-        return self._get_library("stresses").list_symbols()
-
-    @property
-    def model_names(self):
-        """List of model names.
-
-        Returns
-        -------
-        list
-            list of models in library
-        """
-        return self._get_library("models").list_symbols()
-
-    @property
-    def oseries_with_models(self):
-        """List of oseries with models."""
-        return self._get_library("oseries_models").list_symbols()
-
-
 class ArcticDBConnector(BaseConnector, ConnectorUtil):
+    """ArcticDBConnector object using ArcticDB to store data."""
+
     conn_type = "arcticdb"
 
     def __init__(self, name: str, uri: str):
@@ -228,8 +48,7 @@ class ArcticDBConnector(BaseConnector, ConnectorUtil):
         self._update_all_oseries_model_links()
 
     def _initialize(self) -> None:
-        """Internal method to initalize the libraries."""
-
+        """Initialize the libraries (internal method)."""
         for libname in self._default_library_names:
             if self._library_name(libname) not in self.arc.list_libraries():
                 self.arc.create_library(self._library_name(libname))
@@ -242,7 +61,7 @@ class ArcticDBConnector(BaseConnector, ConnectorUtil):
             self.libs[libname] = self._get_library(libname)
 
     def _library_name(self, libname: str) -> str:
-        """Internal method to get full library name according to ArcticDB."""
+        """Get full library name according to ArcticDB (internal method)."""
         return ".".join([self.name, libname])
 
     def _get_library(self, libname: str):
@@ -270,7 +89,7 @@ class ArcticDBConnector(BaseConnector, ConnectorUtil):
         metadata: Optional[Dict] = None,
         **_,
     ) -> None:
-        """Internal method to add item to library (time series or model).
+        """Add item to library (time series or model) (internal method).
 
         Parameters
         ----------
@@ -292,7 +111,7 @@ class ArcticDBConnector(BaseConnector, ConnectorUtil):
             lib.write(name, item, metadata=metadata)
 
     def _get_item(self, libname: str, name: str) -> Union[FrameorSeriesUnion, Dict]:
-        """Internal method to retrieve item from library.
+        """Retrieve item from library (internal method).
 
         Parameters
         ----------
@@ -310,7 +129,7 @@ class ArcticDBConnector(BaseConnector, ConnectorUtil):
         return lib.read(name).data
 
     def _del_item(self, libname: str, name: str) -> None:
-        """Internal method to delete items (series or models).
+        """Delete items (series or models) (internal method).
 
         Parameters
         ----------
@@ -323,7 +142,7 @@ class ArcticDBConnector(BaseConnector, ConnectorUtil):
         lib.delete(name)
 
     def _get_metadata(self, libname: str, name: str) -> dict:
-        """Internal method to retrieve metadata for an item.
+        """Retrieve metadata for an item (internal method).
 
         Parameters
         ----------
@@ -379,227 +198,9 @@ class ArcticDBConnector(BaseConnector, ConnectorUtil):
         return self._get_library("oseries_models").list_symbols()
 
 
-class PystoreConnector(BaseConnector, ConnectorUtil):  # pragma: no cover
-    conn_type = "pystore"
-
-    def __init__(self, name: str, path: str):
-        """Create a PystoreConnector object that points to a Pystore.
-
-        Parameters
-        ----------
-        name : str
-            name of the store
-        path : str
-            path to the pystore directory
-        """
-        try:
-            import pystore
-        except ModuleNotFoundError as e:
-            print(
-                "Install pystore, follow instructions at "
-                "https://github.com/ranaroussi/pystore#dependencies"
-            )
-            raise e
-        self.name = name
-        self.path = path
-        pystore.set_path(self.path)
-        self.store = pystore.store(self.name)
-        self.libs: dict = {}
-        self._initialize()
-        self.models = ModelAccessor(self)
-        # for older versions of PastaStore, if oseries_models library is empty
-        # populate oseries - models database
-        self._update_all_oseries_model_links()
-
-    def _initialize(self) -> None:
-        """Internal method to initalize the libraries (stores)."""
-        for libname in self._default_library_names:
-            if libname in self.store.list_collections():
-                print(
-                    f"PystoreConnector: library '{self.path}/{libname}' "
-                    "already exists. Linking to existing library."
-                )
-            lib = self.store.collection(libname)
-            self.libs[libname] = lib
-
-    def _get_library(self, libname: str):
-        """Get Pystore library handle.
-
-        Parameters
-        ----------
-        libname : str
-            name of the library
-
-        Returns
-        -------
-        Pystore.Collection handle
-            handle to the library
-        """
-        # get library handle
-        lib = self.store.collection(libname)
-        return lib
-
-    def _add_item(
-        self,
-        libname: str,
-        item: Union[FrameorSeriesUnion, Dict],
-        name: str,
-        metadata: Optional[Dict] = None,
-        overwrite: bool = False,
-    ) -> None:
-        """Internal method to add item to library (time series or model).
-
-        Parameters
-        ----------
-        libname : str
-            name of the library
-        item : Union[FrameorSeriesUnion, Dict]
-            item to add, either time series or pastas.Model as dictionary
-        name : str
-            name of the item
-        metadata : Optional[Dict], optional
-            dictionary containing metadata, by default None
-        overwrite : bool, optional
-            overwrite item if it already exists, by default False.
-        """
-        # convert to DataFrame because pystore doesn't accept pandas.Series
-        # (maybe has an easy fix, but converting w to_frame for now)
-        if isinstance(item, pd.Series):
-            s = item.to_frame(name=name)
-            is_type = "series"
-        elif isinstance(item, dict):
-            s = pd.DataFrame()  # empty DataFrame as placeholder
-            jsondict = json.loads(json.dumps(item, cls=PastasEncoder, indent=4))
-            metadata = jsondict  # model dict is stored in metadata
-            is_type = "series"
-        elif isinstance(item, list):
-            s = pd.Series(item).to_frame(name="modelnames")
-            is_type = "list"
-        elif isinstance(item, pd.DataFrame):
-            s = item
-            is_type = "dataframe"
-
-        # store info about input type to ensure same type is returned
-        if metadata is None:
-            metadata = {"_is_type": is_type}
-        else:
-            metadata["_is_type"] = is_type
-
-        lib = self._get_library(libname)
-        lib.write(name, s, metadata=metadata, overwrite=overwrite)
-
-    def _get_item(self, libname: str, name: str) -> Union[FrameorSeriesUnion, Dict]:
-        """Internal method to retrieve item from pystore library.
-
-        Parameters
-        ----------
-        libname : str
-            name of the library
-        name : str
-            name of the item
-
-        Returns
-        -------
-        item : Union[FrameorSeriesUnion, Dict]
-            time series or model dictionary
-        """
-        load_mod = import_module("pastas.io.pas")  # type: ignore
-        lib = self._get_library(libname)
-        # hack for storing models, stored as metadata
-        if libname == "models":
-            jsonpath = lib._item_path(name).joinpath("metadata.json")
-            s = load_mod.load(jsonpath)  # type: ignore
-        else:
-            # read series and convert to pandas
-            item = lib.item(name)
-            s = item.to_pandas()
-            # remove _is_type key and return correct type
-            is_type = item.metadata.pop("_is_type")
-            if is_type == "series":
-                s = s.squeeze()
-            elif is_type == "list":
-                s = s["modelnames"].tolist()
-        return s
-
-    def _del_item(self, libname: str, name: str) -> None:
-        """Internal method to delete data from the store.
-
-        Parameters
-        ----------
-        libname : str
-            name of the library
-        name : str
-            name of the item to delete
-        """
-        lib = self._get_library(libname)
-        lib.delete_item(name)
-        self._clear_cache(libname)
-
-    def _get_metadata(self, libname: str, name: str) -> dict:
-        """Internal method to read metadata from pystore.
-
-        Parameters
-        ----------
-        libname : str
-            name of the library the series are in ("oseries" or "stresses")
-        name : str
-            name of item to load metadata for
-
-        Returns
-        -------
-        imeta : dict
-            dictionary containing metadata
-        """
-        from pystore.utils import read_metadata
-
-        lib = self._get_library(libname)
-        imeta = read_metadata(lib._item_path(name))
-        if "name" not in imeta.keys():
-            imeta["name"] = name
-        if "_is_type" in imeta.keys():
-            imeta.pop("_is_type")
-        return imeta
-
-    @property
-    def oseries_names(self):
-        """List of oseries names.
-
-        Returns
-        -------
-        list
-            list of oseries in library
-        """
-        return list(self._get_library("oseries").list_items())
-
-    @property
-    def stresses_names(self):
-        """List of stresses names.
-
-        Returns
-        -------
-        list
-            list of stresses in library
-        """
-        return list(self._get_library("stresses").list_items())
-
-    @property
-    def model_names(self):
-        """List of model names.
-
-        Returns
-        -------
-        list
-            list of models in library
-        """
-        return list(self._get_library("models").list_items())
-
-    @property
-    def oseries_with_models(self):
-        """List of oseries with models."""
-        return list(self._get_library("oseries_models").list_items())
-
-
 class DictConnector(BaseConnector, ConnectorUtil):
+    """DictConnector object that stores timeseries and models in dictionaries."""
+
     conn_type = "dict"
 
     def __init__(self, name: str = "pastas_db"):
@@ -643,7 +244,7 @@ class DictConnector(BaseConnector, ConnectorUtil):
         metadata: Optional[Dict] = None,
         **_,
     ) -> None:
-        """Internal method to add item (time series or models).
+        """Add item (time series or models) (internal method).
 
         Parameters
         ----------
@@ -663,7 +264,7 @@ class DictConnector(BaseConnector, ConnectorUtil):
             lib[name] = (metadata, item)
 
     def _get_item(self, libname: str, name: str) -> Union[FrameorSeriesUnion, Dict]:
-        """Internal method to retrieve item from pystore library.
+        """Retrieve item from database (internal method).
 
         Parameters
         ----------
@@ -685,7 +286,7 @@ class DictConnector(BaseConnector, ConnectorUtil):
         return item
 
     def _del_item(self, libname: str, name: str) -> None:
-        """Internal method to delete items (series or models).
+        """Delete items (series or models) (internal method).
 
         Parameters
         ----------
@@ -698,7 +299,7 @@ class DictConnector(BaseConnector, ConnectorUtil):
         _ = lib.pop(name)
 
     def _get_metadata(self, libname: str, name: str) -> dict:
-        """Internal method to read metadata.
+        """Read metadata (internal method).
 
         Parameters
         ----------
@@ -742,6 +343,8 @@ class DictConnector(BaseConnector, ConnectorUtil):
 
 
 class PasConnector(BaseConnector, ConnectorUtil):
+    """PasConnector object that stores time series and models as JSON files on disk."""
+
     conn_type = "pas"
 
     def __init__(self, name: str, path: str):
@@ -767,7 +370,7 @@ class PasConnector(BaseConnector, ConnectorUtil):
         self._update_all_oseries_model_links()
 
     def _initialize(self) -> None:
-        """Internal method to initialize the libraries."""
+        """Initialize the libraries (internal method)."""
         for val in self._default_library_names:
             libdir = os.path.join(self.path, val)
             if not os.path.exists(libdir):
@@ -803,7 +406,7 @@ class PasConnector(BaseConnector, ConnectorUtil):
         metadata: Optional[Dict] = None,
         **_,
     ) -> None:
-        """Internal method to add item (time series or models).
+        """Add item (time series or models) (internal method).
 
         Parameters
         ----------
@@ -845,7 +448,7 @@ class PasConnector(BaseConnector, ConnectorUtil):
                 fm.write(jsondict)
 
     def _get_item(self, libname: str, name: str) -> Union[FrameorSeriesUnion, Dict]:
-        """Internal method to retrieve item.
+        """Retrieve item (internal method).
 
         Parameters
         ----------
@@ -878,7 +481,7 @@ class PasConnector(BaseConnector, ConnectorUtil):
         return item
 
     def _del_item(self, libname: str, name: str) -> None:
-        """Internal method to delete items (series or models).
+        """Delete items (series or models) (internal method).
 
         Parameters
         ----------
@@ -898,7 +501,7 @@ class PasConnector(BaseConnector, ConnectorUtil):
                 pass
 
     def _get_metadata(self, libname: str, name: str) -> dict:
-        """Internal method to read metadata.
+        """Read metadata (internal method).
 
         Parameters
         ----------

pastastore/datasets.py

@@ -1,3 +1,5 @@
+"""Module containing example dataset."""
+
 import os
 
 import pandas as pd
@@ -13,7 +15,7 @@ from pastastore.base import BaseConnector
 
 
 def example_pastastore(conn="DictConnector"):
-    """Example dataset loaded into PastaStore.
+    """Get example PastaStore.
 
     Parameters
     ----------
@@ -27,7 +29,6 @@ def example_pastastore(conn="DictConnector"):
     pstore : pastastore.PastaStore
         PastaStore containing example dataset
     """
-
     # check it test dataset is available
     datadir = os.path.join(os.path.dirname(__file__), "../tests/data")
     if not os.path.exists(datadir):
@@ -175,8 +176,7 @@ def _default_connector(conntype: str):
     Parameters
     ----------
     conntype : str
-        name of connector (DictConnector, PasConnector,
-        ArcticConnector, ArcticDBConnector or PystoreConnector)
+        name of connector (DictConnector, PasConnector, ArcticDBConnector)
 
     Returns
     -------
@@ -184,14 +184,9 @@ def _default_connector(conntype: str):
         default Connector based on type.
     """
     Conn = getattr(pst, conntype)
-    if Conn.conn_type == "arctic":
-        connstr = "mongodb://localhost:27017/"
-        conn = Conn("my_db", connstr)
-    elif Conn.conn_type == "arcticdb":
+    if Conn.conn_type == "arcticdb":
         uri = "lmdb://./arctic_db"
         conn = Conn("my_db", uri)
-    elif Conn.conn_type == "pystore":
-        conn = Conn("my_db", "./pystore_db")
     elif Conn.conn_type == "dict":
         conn = Conn("my_db")
     elif Conn.conn_type == "pas":