Repository: cogent3/cogent3
Pull request #1229 (2022-12-14 15:57:04)

move get_data_store to io_new.py and change it to open_data_store


Patch:
diff --git a/src/cogent3/__init__.py b/src/cogent3/__init__.py
index 351ad632a..64bc76398 100644
--- a/src/cogent3/__init__.py
+++ b/src/cogent3/__init__.py
@@ -10,7 +10,7 @@ import warnings
 
 from typing import Callable, Optional, Union
 
-from cogent3.app import app_help, available_apps, get_app
+from cogent3.app import app_help, available_apps, get_app, open_data_store
 from cogent3.core.alignment import (
     Alignment,
     ArrayAlignment,
diff --git a/src/cogent3/app/__init__.py b/src/cogent3/app/__init__.py
index b796dac6a..b41f5bd65 100644
--- a/src/cogent3/app/__init__.py
+++ b/src/cogent3/app/__init__.py
@@ -5,6 +5,8 @@ import importlib
 import inspect
 import re
 
+from .io_new import open_data_store
+
 
 __author__ = "Gavin Huttley"
 __copyright__ = "Copyright 2007-2022, The Cogent Project"
diff --git a/src/cogent3/app/io.py b/src/cogent3/app/io.py
index 986b4998c..bbc5f5319 100644
--- a/src/cogent3/app/io.py
+++ b/src/cogent3/app/io.py
@@ -139,42 +139,20 @@ def findall(base_path, suffix="fa", limit=None, verbose=False):
 def get_data_store(
     base_path: Union[str, pathlib.Path], suffix=None, limit=None, verbose=False
 ):
-    """returns DataStore containing glob matches to suffix in base_path
-
-    Parameters
-    ----------
-    base_path : str or Path
-        path to directory or zipped archive
-    suffix : str
-        suffix of filenames
-    limit : int or None
-        the number of matches to return
-    Returns
-    -------
-    ReadOnlyDirectoryDataStore or ReadOnlyZippedDataStore
-    """
-    base_path = pathlib.Path(base_path)
-    base_path = base_path.expanduser().absolute()
-    if base_path.suffix in (".tinydb", ".sqlitedb"):
-        suffix = "json"
-
-    if suffix is None:
-        raise ValueError("suffix required")
-
-    if not base_path.exists():
-        raise ValueError(f"'{base_path}' does not exist")
-
-    if type(suffix) != str:
-        raise ValueError(f"{suffix} is not a string")
-
-    if zipfile.is_zipfile(base_path):
-        ds_suffix = ".zip"
-    elif base_path.suffix:
-        ds_suffix = base_path.suffix
-    else:
-        ds_suffix = None
-    klass = _datastore_reader_map[ds_suffix]
-    return klass(base_path, suffix=suffix, limit=limit)
+    """DEPRECATED, use top level open_data_store"""
+    from cogent3 import open_data_store
+    from cogent3.util.warning import deprecated
+
+    deprecated(
+        "function",
+        "get_data_store",
+        "cogent3.open_data_store",
+        "2023.3",
+        "renamed and now a top-level import",
+        stack_level=1,
+    )
+
+    return open_data_store(base_path, suffix, limit, verbose)
 
 
 def _load_seqs(path, klass, parser, moltype):
diff --git a/src/cogent3/app/io_new.py b/src/cogent3/app/io_new.py
index d4c189fd9..7ecc7954c 100644
--- a/src/cogent3/app/io_new.py
+++ b/src/cogent3/app/io_new.py
@@ -1,10 +1,13 @@
 import contextlib
 import json
+import os
 import pickle
+import zipfile
 
 from enum import Enum
 from gzip import compress as gzip_compress
 from gzip import decompress as gzip_decompress
+from pathlib import Path
 from typing import Optional, Union
 
 import numpy
@@ -23,6 +26,14 @@ from cogent3.util.deserialise import deserialise_object
 from cogent3.util.table import Table
 
 from .composable import LOADER, WRITER, NotCompleted, define_app
+from .data_store import (
+    ReadOnlyDirectoryDataStore,
+    ReadOnlyTinyDbDataStore,
+    ReadOnlyZippedDataStore,
+    get_data_source,
+    load_record_from_json,
+    make_record_for_json,
+)
 from .data_store_new import (
     DataStoreABC,
     get_data_source,
@@ -49,6 +60,122 @@ __email__ = "Gavin.Huttley@anu.edu.au"
 __status__ = "Alpha"
 
 
+_datastore_reader_map = {}
+
+
+class register_datastore_reader:
+    """
+    registration decorator for read only data store classes
+
+    The registration key must be a string that of the file format suffix
+    (more than one suffix can be registered at a time).
+
+    Parameters
+    ----------
+    args: str or sequence of str
+        must be unique, a preceding '.' will be added if not already present
+    """
+
+    def __init__(self, *args):
+        args = list(args)
+        for i, suffix in enumerate(args):
+            if suffix is None:
+                assert (
+                    suffix not in _datastore_reader_map
+                ), f"{suffix!r} already in {list(_datastore_reader_map)}"
+                continue
+
+            if not isinstance(suffix, str):
+                raise TypeError(f"{suffix!r} is not a string")
+
+            if suffix.strip() == suffix and not suffix:
+                raise ValueError("cannot have white-space suffix")
+
+            suffix = suffix.strip()
+            if suffix:
+                suffix = suffix if suffix[0] == "." else f".{suffix}"
+
+            assert (
+                suffix not in _datastore_reader_map
+            ), f"{suffix!r} already in {list(_datastore_reader_map)}"
+            args[i] = suffix
+
+        self._type_str = tuple(args)
+
+    def __call__(self, func):
+        for type_str in self._type_str:
+            _datastore_reader_map[type_str] = func
+        return func
+
+
+# register the main readers
+register_datastore_reader("zip")(ReadOnlyZippedDataStore)
+register_datastore_reader("tinydb")(ReadOnlyTinyDbDataStore)
+register_datastore_reader(None)(ReadOnlyDirectoryDataStore)
+
+
+def findall(base_path, suffix="fa", limit=None, verbose=False):
+    """returns glob match to suffix, path is relative to base_path
+
+    Parameters
+    ----------
+    base_path : str
+        path to directory or zipped archive
+    suffix : str
+        suffix of filenames
+    limit : int or None
+        the number of matches to return
+    """
+    if not os.path.exists(base_path):
+        raise ValueError(f"'{base_path}' does not exist")
+
+    zipped = zipfile.is_zipfile(base_path)
+    klass = _datastore_reader_map.get(".zip" if zipped else None)
+    data_store = klass(base_path, suffix=suffix, limit=limit, verbose=verbose)
+    return data_store.members
+
+
+def open_data_store(
+    base_path: Union[str, Path], suffix=None, limit=None, verbose=False
+):
+    """returns DataStore containing glob matches to suffix in base_path
+
+    Parameters
+    ----------
+    base_path : str or Path
+        path to directory or zipped archive
+    suffix : str
+        suffix of filenames
+    limit : int or None
+        the number of matches to return
+    Returns
+    -------
+    ReadOnlyDirectoryDataStore or ReadOnlyZippedDataStore
+    """
+    base_path = Path(base_path)
+    base_path = base_path.expanduser().absolute()
+    if base_path.suffix in (".tinydb", ".sqlitedb"):
+        suffix = "json"
+
+    if suffix is None:
+        raise ValueError("suffix required")
+
+    if not base_path.exists():
+        raise ValueError(f"'{base_path}' does not exist")
+
+    if type(suffix) != str:
+        raise ValueError(f"{suffix} is not a string")
+
+    if zipfile.is_zipfile(base_path):
+        ds_suffix = ".zip"
+    elif base_path.suffix:
+        ds_suffix = base_path.suffix
+    else:
+        ds_suffix = None
+    klass = _datastore_reader_map[ds_suffix]
+    return klass(base_path, suffix=suffix, limit=limit)
+
+
 @define_app
 def pickle_it(data: SerialisableType) -> bytes:
     return pickle.dumps(data)
diff --git a/src/cogent3/evolve/fast_distance.py b/src/cogent3/evolve/fast_distance.py
index 55ecb1e63..35e6e789e 100644
--- a/src/cogent3/evolve/fast_distance.py
+++ b/src/cogent3/evolve/fast_distance.py
@@ -6,7 +6,7 @@ import numpy
 from numpy import array, diag, dot, eye, float64, int32, log, sqrt, zeros
 from numpy.linalg import det, inv
 
-from cogent3 import DNA, RNA, get_moltype
+from cogent3.core.moltype import DNA, RNA, get_moltype
 from cogent3.util.dict_array import DictArray
 from cogent3.util.misc import get_object_provenance
 from cogent3.util.progress_display import display_wrap


Test patch:
diff --git a/tests/test_app/test_app_mpi.py b/tests/test_app/test_app_mpi.py
index d5c8e9b49..d5480df6a 100644
--- a/tests/test_app/test_app_mpi.py
+++ b/tests/test_app/test_app_mpi.py
@@ -3,7 +3,8 @@ from tempfile import TemporaryDirectory
 from unittest import TestCase, main, skipUnless
 
 from cogent3.app import align as align_app
-from cogent3.app import io as io_app
+from cogent3.app import io_new as io_app
+from cogent3.app import open_data_store
 from cogent3.app.io import write_db
 from cogent3.util import parallel
 
@@ -24,7 +25,7 @@ class MPITests(TestCase):
     @skipUnless(parallel.USING_MPI, reason="Not using MPI")
     def test_write_db(self):
         """writing with overwrite in MPI should reset db"""
-        dstore = io_app.get_data_store("data", suffix="fasta")
+        dstore = open_data_store("data", suffix="fasta")
         members = dstore.filtered(callback=lambda x: "brca1.fasta" not in x.split("/"))
         with TemporaryDirectory(dir=".") as dirname:
             path = Path(dirname) / "delme.tinydb"
@@ -44,7 +45,7 @@ class MPITests(TestCase):
             process.data_store.close()
 
             # now get read only and check what's in there
-            result = io_app.get_data_store(path)
+            result = open_data_store(path)
             got = [str(m) for m in result]
 
             assert got == expect
diff --git a/tests/test_app/test_composable.py b/tests/test_app/test_composable.py
index 4e7245ae1..25ad01ddf 100644
--- a/tests/test_app/test_composable.py
+++ b/tests/test_app/test_composable.py
@@ -13,7 +13,7 @@ import pytest
 from numpy import array, ndarray
 from scitrack import CachingLogger
 
-from cogent3 import make_aligned_seqs
+from cogent3 import make_aligned_seqs, open_data_store
 from cogent3.app import align, evo
 from cogent3.app import io as io_app
 from cogent3.app import io_new as io_app_new
@@ -303,7 +303,7 @@ def test_disconnect():
 
 def test_as_completed():
     """correctly applies iteratively"""
-    dstore = io_app.get_data_store(DATA_DIR, suffix="fasta", limit=3)
+    dstore = open_data_store(DATA_DIR, suffix="fasta", limit=3)
     reader = io_app.load_unaligned(format="fasta", moltype="dna")
     got = list(reader.as_completed(dstore))
     assert len(got) == len(dstore)
@@ -334,7 +334,7 @@ def test_as_completed():
 @pytest.mark.parametrize("klass", (DataStoreDirectory,))
 def test_apply_to_strings(tmp_dir, klass):
     """apply_to handles strings as paths"""
-    dstore = io_app.get_data_store(DATA_DIR, suffix="fasta", limit=3)
+    dstore = open_data_store(DATA_DIR, suffix="fasta", limit=3)
     dstore = [str(m) for m in dstore]
     reader = io_app.load_aligned(format="fasta", moltype="dna")
     min_length = sample_app.min_length(10)
@@ -365,7 +365,7 @@ def test_apply_to_non_unique_identifiers(tmp_dir):
 
 def test_apply_to_logging(tmp_dir):
     """correctly creates log file"""
-    dstore = io_app.get_data_store(DATA_DIR, suffix="fasta", limit=3)
+    dstore = open_data_store(DATA_DIR, suffix="fasta", limit=3)
     reader = io_app.load_aligned(format="fasta", moltype="dna")
     min_length = sample_app.min_length(10)
     outpath = os.path.join(os.getcwd(), tmp_dir, "delme.tinydb")
@@ -379,7 +379,7 @@ def test_apply_to_logging(tmp_dir):
 
 def test_apply_to_logger(tmp_dir):
     """correctly uses user provided logger"""
-    dstore = io_app.get_data_store(DATA_DIR, suffix="fasta", limit=3)
+    dstore = open_data_store(DATA_DIR, suffix="fasta", limit=3)
     LOGGER = CachingLogger()
     reader = io_app.load_aligned(format="fasta", moltype="dna")
     min_length = sample_app.min_length(10)
@@ -393,7 +393,7 @@ def test_apply_to_logger(tmp_dir):
 
 def test_apply_to_invalid_logger(tmp_dir):
     """incorrect logger value raises TypeError"""
-    dstore = io_app.get_data_store(DATA_DIR, suffix="fasta", limit=3)
+    dstore = open_data_store(DATA_DIR, suffix="fasta", limit=3)
     for logger_val in (True, "somepath.log"):
         reader = io_app.load_aligned(format="fasta", moltype="dna")
         min_length = sample_app.min_length(10)
@@ -407,7 +407,7 @@ def test_apply_to_invalid_logger(tmp_dir):
 
 def test_apply_to_not_completed(tmp_dir):
     """correctly creates notcompleted"""
-    dstore = io_app.get_data_store(DATA_DIR, suffix="fasta", limit=3)
+    dstore = open_data_store(DATA_DIR, suffix="fasta", limit=3)
     reader = io_app.load_aligned(format="fasta", moltype="dna")
     # trigger creation of notcompleted
     min_length = sample_app.min_length(3000)
@@ -421,7 +421,7 @@ def test_apply_to_not_completed(tmp_dir):
 
 def test_apply_to_not_partially_done(tmp_dir):
     """correctly applies process when result already partially done"""
-    dstore = io_app.get_data_store(DATA_DIR, suffix="fasta")
+    dstore = open_data_store(DATA_DIR, suffix="fasta")
     num_records = len(dstore)
     dirname = Path(tmp_dir)
     reader = io_app.load_aligned(format="fasta", moltype="dna")
@@ -434,7 +434,7 @@ def test_apply_to_not_partially_done(tmp_dir):
     process = reader + writer
     _ = process.apply_to(dstore, show_progress=False)
     writer.data_store.close()
-    dstore = io_app.get_data_store(outpath)
+    dstore = open_data_store(outpath)
     assert len(dstore) == num_records
     dstore.close()
 
@@ -1116,7 +1116,7 @@ def test_apply_to_only_appends(half_dstore1, half_dstore2):
     writer1 = io_app_new.write_seqs(half_dstore1)
     process1 = reader1 + min_length1 + writer1
     # create paths as strings
-    dstore1 = io_app.get_data_store(half_dstore1.source, suffix="fasta")
+    dstore1 = open_data_store(half_dstore1.source, suffix="fasta")
     dstore1 = [str(m) for m in dstore1]
     # check fail on append the same records
     with pytest.raises(IOError):
diff --git a/tests/test_app/test_io.py b/tests/test_app/test_io.py
index 2768f6e74..4ba383e09 100644
--- a/tests/test_app/test_io.py
+++ b/tests/test_app/test_io.py
@@ -13,7 +13,7 @@ import pytest
 
 from numpy.testing import assert_allclose
 
-from cogent3 import DNA
+from cogent3 import DNA, open_data_store
 from cogent3.app import align as align_app
 from cogent3.app import io as io_app
 from cogent3.app.composable import NotCompleted
@@ -52,53 +52,27 @@ class TestIo(TestCase):
         if path.exists():
             path.unlink()
 
-    def test_findall(self):
-        """find all files recursively"""
-        found = list(io_app.findall(self.basedir, suffix=".fasta"))
-        self.assertTrue(len(found) > 1)
-        found = list(io_app.findall(self.basedir, suffix=".fasta", limit=2))
-        self.assertTrue(len(found) == 2)
-
-        # and with a suffix
-        found = list(io_app.findall(self.basedir, suffix=".fasta*"))
-        self.assertTrue(len(found) > 2)
-
-    def test_findall_zip(self):
-        """find all files recursively in a zip archive"""
-        with TemporaryDirectory(dir=".") as dirname:
-            zip_path = join(dirname, "new")
-            shutil.make_archive(zip_path, "zip", self.basedir)
-            zip_path += ".zip"  # because shutil adds the suffix
-            found = list(io_app.findall(zip_path, suffix="fasta"))
-            self.assertTrue(len(found) > 1)
-            found = list(io_app.findall(zip_path, suffix="fasta", limit=2))
-            self.assertTrue(len(found) == 2)
-
-            # and with a suffix
-            found = list(io_app.findall(zip_path, suffix=".fasta*"))
-            self.assertTrue(len(found) > 2)
-
     def test_define_data_store(self):
         """returns an iterable data store"""
-        found = io_app.get_data_store(self.basedir, suffix=".fasta")
+        found = open_data_store(self.basedir, suffix=".fasta")
         self.assertTrue(len(found) > 1)
-        found = io_app.get_data_store(self.basedir, suffix=".fasta", limit=2)
+        found = open_data_store(self.basedir, suffix=".fasta", limit=2)
         self.assertTrue(len(found) == 2)
 
         # and with a suffix
-        found = list(io_app.get_data_store(self.basedir, suffix=".fasta*"))
+        found = list(open_data_store(self.basedir, suffix=".fasta*"))
         self.assertTrue(len(found) > 2)
 
         # with a wild-card suffix
-        found = list(io_app.get_data_store(self.basedir, suffix="*"))
+        found = list(open_data_store(self.basedir, suffix="*"))
         self.assertEqual(len(os.listdir(self.basedir)), len(found))
 
         # raises ValueError if suffix not provided or invalid
         with self.assertRaises(ValueError):
-            _ = io_app.get_data_store(self.basedir)
+            _ = open_data_store(self.basedir)
 
         with self.assertRaises(ValueError):
-            _ = io_app.get_data_store(self.basedir, 1)
+            _ = open_data_store(self.basedir, 1)
 
     def test_load_aligned(self):
         """correctly loads aligned seqs"""
@@ -110,13 +84,13 @@ class TestIo(TestCase):
                 self.assertIsInstance(aln, ArrayAlignment)
                 self.assertEqual(aln.info.source, paths[i])
 
-        fasta_paths = io_app.get_data_store(self.basedir, suffix=".fasta", limit=2)
+        fasta_paths = open_data_store(self.basedir, suffix=".fasta", limit=2)
         fasta_loader = io_app.load_aligned(format="fasta")
         validate(fasta_paths, fasta_loader)
 
     def test_load_aligned_nexus(self):
         """should handle nexus too"""
-        nexus_paths = io_app.get_data_store(self.basedir, suffix="nex")
+        nexus_paths = open_data_store(self.basedir, suffix="nex")
         loader = io_app.load_aligned(format="nexus")
         results = [loader(m) for m in nexus_paths]
         for result in results:
@@ -124,7 +98,7 @@ class TestIo(TestCase):
 
     def test_load_aligned_paml(self):
         """should handle paml too"""
-        paml_paths = io_app.get_data_store(self.basedir, suffix="paml")
+        paml_paths = open_data_store(self.basedir, suffix="paml")
         loader = io_app.load_aligned(format="paml")
         results = [loader(m) for m in paml_paths]
         for result in results:
@@ -155,7 +129,7 @@ class TestIo(TestCase):
 
     def test_load_unaligned(self):
         """load_unaligned returns degapped sequence collections"""
-        fasta_paths = io_app.get_data_store(self.basedir, suffix=".fasta", limit=2)
+        fasta_paths = open_data_store(self.basedir, suffix=".fasta", limit=2)
         fasta_loader = io_app.load_unaligned(format="fasta")
         for i, seqs in enumerate(map(fasta_loader, fasta_paths)):
             self.assertIsInstance(seqs, SequenceCollection)
@@ -203,7 +177,7 @@ class TestIo(TestCase):
             with zipfile.ZipFile(zip_path, "a") as out:
                 out.writestr(outpath, data)
 
-            dstore = io_app.get_data_store(zip_path, suffix="json")
+            dstore = open_data_store(zip_path, suffix="json")
             member = dstore.get_member("delme.json")
             reader = io_app.load_json()
             got = reader(member)
@@ -220,7 +194,7 @@ class TestIo(TestCase):
             got = writer(gr)
             self.assertIsInstance(got, DataStoreMember)
             writer.data_store.db.close()
-            dstore = io_app.get_data_store(f"{outpath}.tinydb", suffix="json")
+            dstore = open_data_store(f"{outpath}.tinydb", suffix="json")
             reader = io_app.load_db()
             got = reader(dstore[0])
             dstore.close()
@@ -236,7 +210,7 @@ class TestIo(TestCase):
             data = dict(a=[1, 2], b="string")
             _ = writer(data, identifier=join("blah", "delme.json"))
             writer.data_store.db.close()
-            dstore = io_app.get_data_store(f"{outpath}.tinydb", suffix="json")
+            dstore = open_data_store(f"{outpath}.tinydb", suffix="json")
             reader = io_app.load_db()
             got = reader(dstore[0])
             dstore.close()
@@ -510,7 +484,7 @@ class TestIo(TestCase):
         """writing with overwrite in parallel should reset db"""
         with TemporaryDirectory(dir=".") as dirname:
             outdir = join(dirname, "delme.tinydb")
-            dstore = io_app.get_data_store(self.basedir, suffix="fasta")
+            dstore = open_data_store(self.basedir, suffix="fasta")
             members = dstore.filtered(
                 callback=lambda x: "brca1.fasta" not in x.split("/")
             )
@@ -526,7 +500,7 @@ class TestIo(TestCase):
             process.data_store.close()
 
             # now get read only and check what's in there
-            result = io_app.get_data_store(outdir)
+            result = open_data_store(outdir)
             got = [str(m) for m in result]
             self.assertNotEqual(got, [])
             result.close()