Repository: cogent3/cogent3
Pull request #1161 (2022-10-04 05:03:51)

Data store refactor


Patch:
diff --git a/src/cogent3/app/data_store_new.py b/src/cogent3/app/data_store_new.py
new file mode 100644
index 000000000..92557fbcb
--- /dev/null
+++ b/src/cogent3/app/data_store_new.py
@@ -0,0 +1,346 @@
+from __future__ import annotations
+
+import glob
+import re
+import shutil
+
+from abc import ABC, abstractmethod
+from enum import Enum
+from functools import singledispatch
+from pathlib import Path
+
+from scitrack import get_text_hexdigest
+
+from cogent3.app.typing import TabularType
+from cogent3.util.io import open_
+from cogent3.util.parallel import is_master_process
+
+
+__author__ = "Gavin Huttley"
+__copyright__ = "Copyright 2007-2022, The Cogent Project"
+__credits__ = ["Gavin Huttley"]
+__license__ = "BSD-3"
+__version__ = "2022.8.24a1"
+__maintainer__ = "Gavin Huttley"
+__email__ = "Gavin.Huttley@anu.edu.au"
+__status__ = "Alpha"
+
+_NOT_COMPLETED_TABLE = "not_completed"
+_LOG_TABLE = "logs"
+_MD5_TABLE = "md5"
+
+
+class IfExist(Enum):
+    SKIP = "skip"
+    OVERWRITE = "overwrite"
+    RAISE = "raise"
+    READONLY = "readonly"
+
+
+SKIP = IfExist.SKIP
+OVERWRITE = IfExist.OVERWRITE
+RAISE = IfExist.RAISE
+READONLY = IfExist.READONLY
+
+
+class DataMemberABC(ABC):
+    def __init__(self, data_store: "DataStoreABC" = None):
+        self.data_store = data_store
+
+    @property
+    @abstractmethod
+    def unique_id(self):
+        ...
+
+    def read(self) -> str | bytes:
+        return self.data_store.read(self.unique_id)
+
+    def write(self, data: str) -> None:
+        self.data_store.write(self.unique_id, data)
+
+
+class DataStoreABC(ABC):
+    def __init__(
+        self,
+        source,
+        if_dest_exists: IfExist = READONLY,
+        if_member_exists: IfExist = RAISE,
+    ):
+        self._if_dest_exists = if_dest_exists
+        self._if_member_exists = if_member_exists
+
+    @abstractmethod
+    def __repr__(self):
+        ...
+
+    def __getitem__(self, index):
+        return self.members[index]
+
+    def __len__(self):
+        return len(self.members)
+
+    def __contains__(self, identifier):
+        """whether relative identifier has been stored"""
+        return any(identifier == m.unique_id for m in self)
+
+    @abstractmethod
+    def read(self, unique_id: str) -> str | bytes:
+        ...
+
+    @abstractmethod
+    def write(self, unique_id: str, data: str | bytes) -> None:
+        if self._if_dest_exists is READONLY:
+            raise IOError("datastore is readonly")
+        if self._if_member_exists is RAISE and unique_id in self:
+            raise IOError("member exists")
+
+    @abstractmethod
+    def write_not_completed(self, unique_id: str, data: str | bytes) -> None:
+        if self.if_member_exists is RAISE and unique_id in self:
+            raise IOError("member exists")
+
+    @abstractmethod
+    def write_log(self, unique_id: str, data: str | bytes) -> None:
+        if self.if_member_exists is RAISE and unique_id in self:
+            raise IOError("member exists")
+
+    @abstractmethod
+    def describe(self) -> TabularType:
+        ...
+
+    @property
+    @abstractmethod
+    def members(self) -> list[DataMemberABC]:
+        ...
+
+    def __iter__(self):
+        yield from self.members
+
+    @property
+    @abstractmethod
+    def logs(self) -> list[DataMemberABC]:
+        ...
+
+    @property
+    @abstractmethod
+    def completed(self) -> list[DataMemberABC]:
+        ...
+
+    @property
+    @abstractmethod
+    def not_completed(self) -> list[DataMemberABC]:
+        ...
+
+    @property
+    @abstractmethod
+    def summary_logs(self) -> TabularType:
+        ...
+
+    @property
+    @abstractmethod
+    def summary_not_completed(self) -> TabularType:
+        ...
+
+
+class DataMember(DataMemberABC):
+    def __init__(self, data_store: DataStoreABC = None, unique_id: str = None):
+        super().__init__(data_store)
+        self._unique_id = unique_id
+
+    def __repr__(self):
+        return self._unique_id
+
+    @property
+    def unique_id(self):
+        return self._unique_id
+
+
+class DataStoreDirectory(DataStoreABC):
+    def __init__(
+        self,
+        source,
+        if_dest_exists: IfExist = READONLY,
+        if_member_exists: IfExist = RAISE,
+        suffix=None,
+        limit=None,
+        verbose=False,
+        md5=True,
+    ):
+        super().__init__(source, if_dest_exists, if_member_exists)
+        suffix = suffix or ""
+        if suffix != "*":  # wild card search for all
+            suffix = re.sub(r"^[\s.*]+", "", suffix)  # tidy the suffix
+        source = Path(source)
+        self.source = source.expanduser()
+        self.suffix = suffix
+        self._members = []
+        self._limit = limit
+        self._verbose = verbose
+        self._checksums = {}
+        self._md5 = md5
+        self._source_check_create(if_dest_exists)
+
+    def _source_check_create(self, if_dest_exists):
+        if not is_master_process():
+            return
+
+        sub_dirs = [_NOT_COMPLETED_TABLE, _LOG_TABLE, _MD5_TABLE]
+
+        if if_dest_exists is READONLY and not self.source.exists():
+            raise IOError("must exist")
+        elif if_dest_exists is RAISE and self.source.exists():
+            raise IOError("must not exist")
+        elif if_dest_exists is READONLY:
+            return
+
+        path = Path(self.source)
+        if path.exists() and if_dest_exists == RAISE:
+            raise FileExistsError(f"'{self.source}' exists")
+        elif path.exists() and if_dest_exists == OVERWRITE:
+            for sub_dir in sub_dirs:
+                try:
+                    shutil.rmtree(self.source / sub_dir)
+                except:
+                    pass
+        elif not path.exists():
+            path.mkdir(parents=True, exist_ok=True)
+
+        for sub_dir in sub_dirs:
+            (self.source.parent / sub_dir).mkdir(parents=True, exist_ok=True)
+
+    def read(self, unique_id: str) -> str:
+        """reads data corresponding to identifier"""
+        # before open should make absolute path
+        with open_(self.source / unique_id) as input:
+            data = input.read()
+            if self._md5 and isinstance(data, str):
+                self._checksums[unique_id] = get_text_hexdigest(data)
+
+        return data
+
+    def md5(self, unique_id: str, force=True) -> str:
+        """
+        Parameters
+        ----------
+        identifier
+            name of data store member
+        force : bool
+            forces reading of data if not already done
+
+        Returns
+        -------
+        md5 checksum for the member, if available, None otherwise
+        """
+        md5_setting = self._md5  # for restoring automatic md5 calc setting
+        if force and unique_id not in self._checksums:
+            self._md5 = True
+            _ = self.read(unique_id)
+
+        self._md5 = md5_setting
+        return self._checksums.get(unique_id, None)
+
+    @property
+    def members(self) -> list[DataMember]:
+        if not self._members:  # members in completed and incomplete
+            pattern = f"{self.source}/**/*.{self.suffix}"
+            paths = list(glob.iglob(pattern, recursive=True))
+            self._members = []
+            for i, path in enumerate(paths):
+                if self._limit and i >= self._limit:
+                    break
+                self._members.append(DataMember(self, Path(path).name))
+
+        return self._members
+
+    @property
+    def logs(self) -> list[DataMember]:
+        log_dir = self.source / _LOG_TABLE
+        return (
+            [DataMember(self, Path(_LOG_TABLE) / m.name) for m in log_dir.glob("*")]
+            if log_dir.exists()
+            else []
+        )
+
+    @property
+    def completed(self) -> list[DataMember]:
+        path = Path(self.source)
+        return (
+            [DataMember(self, path / m.name) for m in path.glob(f"*.{self.suffix}")]
+            if path.exists()
+            else []
+        )
+
+    @property
+    def not_completed(self) -> list[DataMember]:
+        ic_dir = self.source / _NOT_COMPLETED_TABLE
+        return (
+            [
+                DataMember(self, Path(_NOT_COMPLETED_TABLE) / m.name)
+                for m in ic_dir.glob("*.json")
+            ]
+            if ic_dir.exists()
+            else []
+        )
+
+    def write(self, unique_id: str, data: str) -> None:
+        super().write(unique_id, data)
+        with open_(self.source / unique_id, mode="w") as out:
+            out.write(data)
+            if self._md5 and isinstance(data, str):
+                self._checksums[unique_id] = get_text_hexdigest(data)
+
+    def write_not_completed(self, unique_id: str, data: str) -> None:
+        super().write(unique_id, data)
+        with open_(self.source / _NOT_COMPLETED_TABLE / unique_id, mode="w") as out:
+            out.write(data)
+
+    def write_log(self, unique_id: str, data: str) -> None:
+        super().write(unique_id, data)
+        with open_(self.source / _LOG_TABLE / unique_id, mode="w") as out:
+            out.write(data)
+
+    def describe(self) -> TabularType:
+        ...
+
+    @property
+    def summary_logs(self) -> TabularType:
+        ...
+
+    @property
+    def summary_not_completed(self) -> TabularType:
+        ...
+
+    def __repr__(self):
+        notcompleted_path = Path(self.source / _NOT_COMPLETED_TABLE)
+        if not notcompleted_path.exists():
+            return "This isn't a results directory"
+        num_completed += len(self)
+        num_not_completed = len(list(notcompleted_path.glob("*")))
+        name = self.__class__.__name__
+        sample = f"{str(list(self[:3]))[:-1]}..." if num_completed > 3 else list(self)
+        return f"{num_completed+num_not_completed}x member {name}(source='{self.source}', members={sample})"
+
+
+@singledispatch
+def get_data_source(data) -> str:
+    ...
+
+
+@get_data_source.register
+def get_str_source(data: str):
+    return data
+
+
+@get_data_source.register
+def get_str_source(data: Path):
+    return str(data)
+
+
+@get_data_source.register
+def get_str_source(data: dict):
+    return str(data.get("source", None))
+
+
+@get_data_source.register
+def get_str_source(data: DataMemberABC):
+    return str(data.unique_id)


Test patch:
diff --git a/tests/test_app/test_data_store_new.py b/tests/test_app/test_data_store_new.py
new file mode 100644
index 000000000..1f8df9390
--- /dev/null
+++ b/tests/test_app/test_data_store_new.py
@@ -0,0 +1,281 @@
+from pathlib import Path
+from pickle import dumps, loads
+
+import pytest
+
+from cogent3.app.composable import NotCompleted
+from cogent3.app.data_store_new import (
+    _LOG_TABLE,
+    _NOT_COMPLETED_TABLE,
+    OVERWRITE,
+    READONLY,
+    DataStoreDirectory,
+)
+
+DATA_DIR = Path(__file__).parent.parent / "data"
+
+@pytest.fixture(scope="session")
+def tmp_dir(tmpdir_factory):
+    return tmpdir_factory.mktemp("datastore")
+
+
+@pytest.fixture(autouse=True)
+def workingdir(tmp_dir, monkeypatch):
+    # this set's the working directory for all tests in this module
+    # as a tmp dir
+    monkeypatch.chdir(tmp_dir)
+
+
+@pytest.fixture(scope="session")
+def fasta_dir(tmp_dir):
+    tmp_dir = Path(tmp_dir)
+    filenames = DATA_DIR.glob("*.fasta")
+    fasta_dir = tmp_dir / "fasta"
+    fasta_dir.mkdir(parents=True, exist_ok=True)
+    for fn in filenames:
+        dest = fasta_dir / fn.name
+        dest.write_text(fn.read_text())
+    return fasta_dir
+
+
+@pytest.fixture(scope="session")
+def write_dir(tmp_dir):
+    tmp_dir = Path(tmp_dir)
+    write_dir = tmp_dir / "write"
+    write_dir.mkdir(parents=True, exist_ok=True)
+    return write_dir
+
+
+@pytest.fixture(scope="session")
+def ic_dir(tmp_dir):
+    tmp_dir = Path(tmp_dir)
+    filenames = DATA_DIR.glob("*.fasta")
+    ic_dir = tmp_dir / "ic_dir"
+    ic_dir.mkdir(parents=True, exist_ok=True)
+    for fn in filenames:
+        dest = ic_dir / fn.name
+        dest.write_text(fn.read_text())
+    logs_dir = ic_dir / _LOG_TABLE
+    nc_dir = ic_dir / _NOT_COMPLETED_TABLE
+    logs_dir.mkdir(exist_ok=True)
+    nc_dir.mkdir(exist_ok=True)
+    (logs_dir / "scitrack.log").write_text((DATA_DIR / "scitrack.log").read_text())
+    nc = NotCompleted("FAIL", "dummy", "dummy_message", source="dummy_source")
+    (nc_dir / "nc.json").write_text(nc.to_json())
+    return ic_dir
+
+
+@pytest.fixture(scope="session")
+def ro_dstore(fasta_dir):
+    return DataStoreDirectory(fasta_dir, suffix=".fasta", if_dest_exists=READONLY)
+
+
+@pytest.fixture(scope="session")
+def w_dstore(write_dir):
+    return DataStoreDirectory(write_dir, if_dest_exists=OVERWRITE)
+
+
+@pytest.fixture(scope="session")
+def ic_dstore(ic_dir):
+    return DataStoreDirectory(ic_dir, suffix=".fasta")
+
+
+def test_contains(ro_dstore):
+    """correctly identify when a data store contains a member"""
+    assert "brca1.fasta" in ro_dstore
+
+
+def test_iter(ro_dstore):
+    """DataStore objects allow iteration over members"""
+    members = list(ro_dstore)
+    assert members == ro_dstore.members
+
+
+def test_len(fasta_dir, ro_dstore):
+    """DataStore returns correct len"""
+    expect = len(list(fasta_dir.glob("*.fasta")))
+    assert expect == len(ro_dstore) == len(ro_dstore.members)
+
+
+def test_getitem(ro_dstore):
+    with pytest.raises(IndexError):
+        _ = ro_dstore[len(ro_dstore)]
+
+    last = ro_dstore[-1]
+    first = ro_dstore[0]
+    assert last.unique_id != first.unique_id
+
+
+def test_iterall(fasta_dir, ro_dstore):
+    expect = {fn.name for fn in fasta_dir.glob("*.fasta")}
+    got = {m.unique_id for m in ro_dstore}
+    assert expect == got
+
+
+def test_read(fasta_dir, ro_dstore):
+    """correctly read content"""
+    expect = (fasta_dir / "brca1.fasta").read_text()
+    got = ro_dstore.read("brca1.fasta")
+    assert got == expect
+
+
+def test_pickleable_roundtrip(ro_dstore):
+    """pickling of data stores should be reversible"""
+    re_dstore = loads(dumps(ro_dstore))
+    assert str(ro_dstore) == str(re_dstore)
+    assert ro_dstore[0].read() == re_dstore[0].read()
+
+
+def test_pickleable_member_roundtrip(ro_dstore):
+    """pickling of data store members should be reversible"""
+    re_member = loads(dumps(ro_dstore[0]))
+    data = re_member.read()
+    assert len(data) > 0
+
+
+def test_empty_directory(fasta_dir):
+    dstore = DataStoreDirectory(fasta_dir, suffix=".txt")
+    assert 0 == len(dstore)
+
+
+def test_no_logs(ro_dstore):
+    assert len(ro_dstore.logs) == 0
+
+
+def test_no_not_completed(ro_dstore):
+    assert len(ro_dstore.not_completed) == 0
+
+
+def test_logs(ic_dstore):
+    assert len(ic_dstore.logs) == 1
+    log = ic_dstore.logs[0].read()
+    assert isinstance(log, str)
+
+
+def test_not_completed(ic_dstore):
+    assert len(ic_dstore.not_completed) == 1
+    nc = ic_dstore.not_completed[0].read()
+    assert isinstance(nc, str)
+
+
+def test_write_read_only_datastore(ro_dstore):
+    with pytest.raises(IOError):
+        ro_dstore.write("brca1.fasta", "test data")
+
+
+def test_write_read_only_member(ro_dstore):
+    with pytest.raises(IOError):
+        ro_dstore[0].write("test data")
+
+
+def test_write(fasta_dir, w_dstore):
+    """correctly write content"""
+    expect = Path(fasta_dir / "brca1.fasta").read_text()
+    identifier = "brca1.fasta"
+    w_dstore.write(identifier, expect)
+    got = w_dstore.read(identifier)
+    assert got == expect
+
+
+def test_multi_write(fasta_dir, w_dstore):
+    """correctly write multiple files to data store"""
+    expect_a = Path(fasta_dir / "brca1.fasta").read_text()
+    expect_b = Path(fasta_dir / "primates_brca1.fasta").read_text()
+    identifier_a = "brca1.fasta"
+    identifier_b = "primates_brca1.fasta"
+    w_dstore.write(identifier_a, expect_a)
+    w_dstore.write(identifier_b, expect_b)
+    got_a = w_dstore.read(identifier_a)
+    got_b = w_dstore.read(identifier_b)
+    # check that both bits of data match
+    assert got_a == expect_a
+    assert got_b == expect_b
+
+
+'''
+def test_write_wout_suffix(w_dstore):
+    """appends suffix expected to records"""
+    with pytest.raises(ValueError):
+        w_dstore.write("1", str(dict(a=24, b="some text")))
+
+    w_dstore.write("1.fasta", str(dict(a=24, b="some text")))
+    assert len(w_dstore)== 1
+
+@skipIf(sys.platform.lower() != "darwin", "broken on linux")
+def test_md5_write(write__dir, w_dstore):
+    """tracks md5 sums of written data"""
+    expect = Path( write_dir / "brca1.fasta").read_text()
+
+    identifier = make_identifier(w_dstore, "brca1.fasta", absolute=True)
+    abs_id = w_dstore.write(identifier, expect)
+    md5 = "05a7302479c55c0b5890b50f617c5642"
+    assert w_dstore.md5(abs_id) == md5
+    assert w_dstore[0].md5 == md5
+
+    # does not have md5 if not set
+    identifier = make_identifier(w_dstore, "brca1.fasta", absolute=True)
+    abs_id = w_dstore.write(identifier, expect)
+    got = w_dstore.md5(abs_id, force=False)
+    assert got is None
+    # but if you set force=True, you get it
+    md5 = "05a7302479c55c0b5890b50f617c5642"
+    got = w_dstore.md5(abs_id, force=True)
+    assert got == md5
+
+
+def test_add_file():
+    """correctly add an arbitrarily named file"""
+    data = Path(f"data{os.sep}brca1.fasta").read_text()
+
+    log_path = os.path.join(dirname, "some.log")
+    with open(log_path, "w") as out:
+        out.write("some text")
+
+    path = os.path.join(dirname, basedir)
+    dstore = DataStoreDirectory(path, suffix=".fa", create=True)
+    _ = dstore.write("brca1.fa", data)
+    dstore.add_file(log_path)
+    assert "some.log" in dstore
+    assert os.path.exists(log_path)
+
+    log_path = os.path.join(dirname, "some.log")
+    with open(log_path, "w") as out:
+        out.write("some text")
+
+    path = os.path.join(dirname, basedir)
+    dstore = DataStoreDirectory(path, suffix=".fa", create=True)
+    _ = dstore.write("brca1.fa", data)
+    dstore.add_file(log_path, cleanup=True)
+    assert "some.log" in dstore
+    assert not os.path.exists(log_path)
+
+def test_make_identifier():
+    """correctly construct an identifier for a new member"""
+
+    if dirname.startswith(f".{os.sep}"):
+        dirname = dirname[2:]
+
+    path = os.path.join(dirname, basedir)
+    base_path = path.replace(".zip", "")
+    dstore = DataStoreDirectory(path, suffix=".json", create=True)
+    name = "brca1.fasta"
+    got = make_identifier(dstore, name, absolute = True)
+    expect = os.path.join(base_path, name.replace("fasta", "json"))
+    assert got == expect
+
+    # now using a DataMember
+    member = DataMember(
+        os.path.join(f"blah{os.sep}blah", f"2-{name}"), None
+    )
+    got = make_identifier(dstore, member, absolute=True)
+    expect = os.path.join(base_path, member.name.replace("fasta", "json"))
+    assert got == expect
+
+
+def test_summary_logs():
+    ...
+
+def test_summary_not_completed():
+   ...
+
+'''