Add property entry methods (#61)

src/atomlite/_internal/database.py

@@ -116,7 +116,7 @@ def __init__(
         )
 
     def num_entries(self) -> int:
-        """Get the number of entries in the database.
+        """Get the number of molecular entries in the database.
 
         .. note::
 
@@ -128,6 +128,18 @@ def num_entries(self) -> int:
             "WHERE molecule IS NOT NULL",
         ).fetchone()[0]
 
+    def num_property_entries(self) -> int:
+        """Get the number of property entries in the database.
+
+        .. note::
+
+            This number includes both the commited and
+            uncommited entries.
+        """
+        return self.connection.execute(
+            f"SELECT COUNT(*) FROM {self._molecule_table}"  # noqa: S608
+        ).fetchone()[0]
+
     def add_entries(
         self,
         entries: Entry | collections.abc.Iterable[Entry],
@@ -186,7 +198,7 @@ def update_entries(
         upsert: bool = True,
         commit: bool = True,
     ) -> None:
-        """Update molecules in the database.
+        """Update molecular entries in the database.
 
         Parameters:
             entries (Entry | list[Entry]):
@@ -246,19 +258,41 @@ def update_entries(
             self.connection.commit()
 
     def has_entry(self, key: str) -> bool:
-        """Check if a molecule is present in the database.
+        """Check if a molecular entry is present in the database.
 
         Parameters:
             key: The key of the molecule to check.
 
         Returns:
-            ``True`` if the molecule is present in the database,
+            ``True`` if the entry is present in the database,
             ``False`` otherwise.
         """
         return (
             self.connection.execute(
                 f"SELECT EXISTS(SELECT 1 FROM {self._molecule_table} "  # noqa: S608
-                "WHERE key=? LIMIT 1)",
+                "WHERE key=? "
+                "AND molecule IS NOT NULL "
+                "LIMIT 1)",
+                (key,),
+            ).fetchone()[0]
+            == 1
+        )
+
+    def has_property_entry(self, key: str) -> bool:
+        """Check if a property entry is present in the database.
+
+        Parameters:
+            key: The key of the molecule to check.
+
+        Returns:
+            ``True`` if the entry is present in the database,
+            ``False`` otherwise.
+        """
+        return (
+            self.connection.execute(
+                f"SELECT EXISTS(SELECT 1 FROM {self._molecule_table} "  # noqa: S608
+                "WHERE key=? "
+                "LIMIT 1)",
                 (key,),
             ).fetchone()[0]
             == 1
@@ -349,6 +383,66 @@ def get_entries(
                 properties=json.loads(properties),
             )
 
+    def get_property_entry(self, key: str) -> PropertyEntry | None:
+        """Get a property entry from the database.
+
+        Parameters:
+            key: The key of the molecule to retrieve from the database.
+
+        Returns:
+            The property entry matching `key` or ``None`` if
+            `key` is not present in the database.
+
+        See Also:
+            * :meth:`.get_property_entries`: For retrieving multiple entries.
+        """
+        result = self.connection.execute(
+            f"SELECT key,properties FROM {self._molecule_table} "  # noqa: S608
+            "WHERE key=? LIMIT 1",
+            (key,),
+        ).fetchone()
+        if result is None:
+            return None
+        key, properties = result
+        return PropertyEntry(key=key, properties=json.loads(properties))
+
+    def get_property_entries(
+        self,
+        keys: str | collections.abc.Iterable[str] | None = None,
+    ) -> collections.abc.Iterator[PropertyEntry]:
+        """Get property entries from the database.
+
+        Parameters:
+            keys (str | list[str] | None):
+                The keys of the molecules to whose properties
+                need to be retrieved from the database.
+                If ``None`` all entries will be returned.
+
+        Yields:
+            A property entry matching `keys`.
+
+        See Also:
+            * :meth:`.get_property_entry`: For retrieving a single entry.
+        """
+        if keys is None:
+            for key, properties in self.connection.execute(
+                f"SELECT key,properties FROM {self._molecule_table}"  # noqa: S608
+            ):
+                yield PropertyEntry(key=key, properties=json.loads(properties))
+            return
+
+        if isinstance(keys, str):
+            keys = (keys,)
+
+        keys = tuple(keys)
+        query = ",".join("?" * len(keys))
+        for key, properties in self.connection.execute(
+            f"SELECT key,properties FROM {self._molecule_table} "  # noqa: S608
+            f"WHERE key IN ({query})",
+            keys,
+        ):
+            yield PropertyEntry(key=key, properties=json.loads(properties))
+
     def get_bool_property(self, key: str, path: str) -> bool | None:
         """Get a boolean property of a molecule.
 
@@ -396,7 +490,7 @@ def set_bool_property(
         *,
         commit: bool = True,
     ) -> None:
-        """Set the property of a molecule.
+        """Set a boolean property of a molecule.
 
         Parameters:
             key:
@@ -463,7 +557,7 @@ def set_int_property(
         *,
         commit: bool = True,
     ) -> None:
-        """Set the property of a molecule.
+        """Set an integer property of a molecule.
 
         Parameters:
             key:
@@ -530,7 +624,7 @@ def set_float_property(
         *,
         commit: bool = True,
     ) -> None:
-        """Set the property of a molecule.
+        """Set a float property of a molecule.
 
         Parameters:
             key:
@@ -597,7 +691,7 @@ def set_str_property(
         *,
         commit: bool = True,
     ) -> None:
-        """Set the property of a molecule.
+        """Set a string property of a molecule.
 
         Parameters:
             key:
@@ -747,7 +841,7 @@ def update_properties(
         merge_properties: bool = True,
         commit: bool = True,
     ) -> None:
-        """Update molecular properties.
+        """Update property entries in the database.
 
         Parameters:
             entries (PropertyEntry | list[PropertyEntry]):

tests/test_database.py

@@ -140,6 +140,39 @@ def test_remove_property(database: atomlite.Database) -> None:
     assert database.get_property("first", "$.a.b") is None
 
 
+def test_get_property_entry(database: atomlite.Database) -> None:
+    database.set_property("first", "$.a.b", 12)
+    prop_entry = database.get_property_entry("first")
+    assert prop_entry is not None
+    assert prop_entry.key == "first"
+    assert prop_entry.properties == {"a": {"b": 12}}
+    prop_entry = database.get_property_entry("second")
+    assert prop_entry is None
+
+
+def test_get_property_entries_returns_all_entries(
+    database: atomlite.Database,
+) -> None:
+    database.set_property("first", "$.a.b", 12)
+    database.set_property("second", "$.a.b", 12)
+    prop_entries = list(database.get_property_entries())
+    assert len(prop_entries) == 2  # noqa: PLR2004
+    assert prop_entries[0].key == "first"
+    assert prop_entries[0].properties == {"a": {"b": 12}}
+    assert prop_entries[1].key == "second"
+    assert prop_entries[1].properties == {"a": {"b": 12}}
+
+
+def test_get_property_entries_returns_entries_of_keys(
+    database: atomlite.Database,
+) -> None:
+    database.set_property("first", "$.a.b", 12)
+    database.set_property("second", "$.a.b", 12)
+    (prop_entry,) = list(database.get_property_entries("first"))
+    assert prop_entry.key == "first"
+    assert prop_entry.properties == {"a": {"b": 12}}
+
+
 def test_num_entries(database: atomlite.Database) -> None:
     assert database.num_entries() == 0
     entry = atomlite.Entry.from_rdkit("first", rdkit.MolFromSmiles("C"))
@@ -151,13 +184,41 @@ def test_num_entries(database: atomlite.Database) -> None:
     entry = atomlite.Entry.from_rdkit("third", rdkit.MolFromSmiles("CCC"))
     database.add_entries(entry, commit=False)
     assert database.num_entries() == 3  # noqa: PLR2004
+    database.set_property("fourth", "$.a", 12)
+    assert database.num_entries() == 3  # noqa: PLR2004
+
+
+def test_property_num_entries(database: atomlite.Database) -> None:
+    assert database.num_property_entries() == 0
+    entry = atomlite.Entry.from_rdkit("first", rdkit.MolFromSmiles("C"))
+    database.add_entries(entry)
+    assert database.num_property_entries() == 1
+    entry = atomlite.Entry.from_rdkit("second", rdkit.MolFromSmiles("CC"))
+    database.add_entries(entry)
+    assert database.num_property_entries() == 2  # noqa: PLR2004
+    entry = atomlite.Entry.from_rdkit("third", rdkit.MolFromSmiles("CCC"))
+    database.add_entries(entry, commit=False)
+    assert database.num_property_entries() == 3  # noqa: PLR2004
+    database.set_property("fourth", "$.a", 12)
+    assert database.num_property_entries() == 4  # noqa: PLR2004
 
 
 def test_has_entry(database: atomlite.Database) -> None:
     entry = atomlite.Entry.from_rdkit("first", rdkit.MolFromSmiles("C"))
     database.add_entries(entry)
+    database.set_property("second", "$.a", 12)
     assert database.has_entry("first")
     assert not database.has_entry("second")
+    assert not database.has_entry("third")
+
+
+def test_has_property_entry(database: atomlite.Database) -> None:
+    entry = atomlite.Entry.from_rdkit("first", rdkit.MolFromSmiles("C"))
+    database.add_entries(entry)
+    database.set_property("second", "$.a", 12)
+    assert database.has_property_entry("first")
+    assert database.has_property_entry("second")
+    assert not database.has_property_entry("third")
 
 
 def test_set_property(database: atomlite.Database) -> None: