Merge pull request #81 from MDverse/update-molecule-model

feat: Add sequence and external_identifiers fields in the Molecule model

Author: pierrepo

src/mdverse_scrapers/models/enums.py

@@ -20,3 +20,10 @@ class DatasetSourceName(StrEnum):
     ATLAS = "atlas"
     GPCRMD = "gpcrmd"
     NMRLIPIDS = "nmrlipids"
+
+
+class ExternalDatabaseName(StrEnum):
+    """External database names."""
+
+    PDB = "pdb"
+    UNIPROT = "uniprot"

src/mdverse_scrapers/models/simulation.py

@@ -5,12 +5,39 @@
 
 from pydantic import BaseModel, ConfigDict, Field, StringConstraints, field_validator
 
+from .enums import ExternalDatabaseName
+
 DOI = Annotated[
     str,
     StringConstraints(pattern=r"^10\.\d{4,9}/[\w\-.]+$"),
 ]
 
 
+class ExternalIdentifier(BaseModel):
+    """External database identifier."""
+
+    # Ensure scraped metadata matches the expected schema exactly
+    # and numbers are coerced to strings when needed.
+    model_config = ConfigDict(extra="forbid", coerce_numbers_to_str=True)
+
+    database_name: ExternalDatabaseName = Field(
+        ...,
+        description=(
+            "Name of the external database. "
+            "Allowed values are defined in ExternalDatabaseName enum. "
+            "Examples: PDB, UNIPROT..."
+        ),
+    )
+    identifier: str = Field(
+        ...,
+        min_length=1,
+        description="Identifier in the external database.",
+    )
+    url: str | None = Field(
+        None, min_length=1, description="Direct URL to the identifier into the database"
+    )
+
+
 class Molecule(BaseModel):
     """Molecule in a simulation."""
 
@@ -19,15 +46,20 @@ class Molecule(BaseModel):
 
     name: str = Field(..., description="Name of the molecule.")
     number_of_atoms: int | None = Field(
-        None, ge=0, description="Number of atoms in the molecule, if known."
+        None, ge=0, description="Number of atoms in the molecule."
     )
-    formula: str | None = Field(
-        None, description="Chemical formula of the molecule, if known."
+    formula: str | None = Field(None, description="Chemical formula of the molecule.")
+    sequence: str | None = Field(
+        None, description="Sequence of the molecule for protein and nucleic acid."
     )
     number_of_molecules: int | None = Field(
         None,
         ge=0,
-        description="Number of molecules of this type in the simulation, if known.",
+        description="Number of molecules of this type in the simulation.",
+    )
+    external_identifiers: list[ExternalIdentifier] | None = Field(
+        None,
+        description=("List of external database identifiers for this molecule."),
     )
 
 

tests/models/test_simulation.py

@@ -3,7 +3,9 @@
 import pytest
 from pydantic import ValidationError
 
+from mdverse_scrapers.models.enums import ExternalDatabaseName
 from mdverse_scrapers.models.simulation import (
+    ExternalIdentifier,
     ForceFieldModel,
     Molecule,
     SimulationMetadata,
@@ -61,7 +63,16 @@ def test_structured_fields_creation():
         software=[Software(name="GROMACS", version="2023.1")],
         molecules=[
             Molecule(
-                name="H2O", number_of_atoms=3, formula="H2O", number_of_molecules=100
+                name="H2O",
+                number_of_atoms=3,
+                formula="H2O",
+                number_of_molecules=100,
+                sequence="PEPTIDE",
+                external_identifiers=[
+                    ExternalIdentifier(
+                        database_name=ExternalDatabaseName.PDB, identifier="1ABC"
+                    )
+                ],
             )
         ],
         forcefields_models=[ForceFieldModel(name="AMBER", version="ff14SB")],
@@ -70,17 +81,21 @@ def test_structured_fields_creation():
     assert metadata.molecules[0].number_of_atoms == 3
     assert metadata.molecules[0].number_of_molecules == 100
     assert metadata.forcefields_models[0].version == "ff14SB"
+    assert metadata.molecules[0].sequence == "PEPTIDE"
+    assert (
+        metadata.molecules[0].external_identifiers[0].database_name
+        == ExternalDatabaseName.PDB
+    )
+    assert metadata.molecules[0].external_identifiers[0].identifier == "1ABC"
 
 
 # -------------------------------------------------------------------
-# Test invalid values in structured fields
+# Test invalid fields
 # -------------------------------------------------------------------
-def test_invalid_molecule_number_of_atoms():
-    """Test that molecule number_of_atoms cannot be negative."""
+def test_invalid_fields():
+    """Test with a non-existing fields."""
     with pytest.raises(ValidationError):
-        SimulationMetadata(
-            molecules=[Molecule(name="H2O", number_of_atoms=-1, formula="H2O")]
-        )
+        SimulationMetadata(total_number_of_something=1000)
 
 
 # -------------------------

tests/models/test_simulation_molecule.py

@@ -0,0 +1,93 @@
+"""Tests for the Pydantic molecule model."""
+
+import pytest
+from pydantic import ValidationError
+
+from mdverse_scrapers.models.enums import ExternalDatabaseName
+from mdverse_scrapers.models.simulation import (
+    ExternalIdentifier,
+    Molecule,
+)
+
+
+# -------------------------------------------------------------------
+# Test invalid fields
+# -------------------------------------------------------------------
+def test_invalid_fields():
+    """Test with a non-existing field."""
+    with pytest.raises(ValidationError):
+        Molecule(
+            name="water",
+            dummy_number=1000,  # type: ignore
+        )
+    with pytest.raises(ValidationError):
+        Molecule(
+            name="water",
+            dummy_str="dummy_value",  # type: ignore
+        )
+
+
+# -------------------------------------------------------------------
+# Test invalid values
+# -------------------------------------------------------------------
+def test_invalid_number_of_atoms():
+    """Test that number_of_atoms cannot be negative."""
+    with pytest.raises(ValidationError):
+        Molecule(name="H2O", number_of_atoms=-10, formula="H2O")
+
+
+def test_invalid_number_of_molecules():
+    """Test that number_of_molecules cannot be negative."""
+    with pytest.raises(ValidationError):
+        Molecule(name="H2O", number_of_molecules=-10, formula="H2O")
+
+
+# -------------------------------------------------------------------
+# Test ExternalIdentifier
+# -------------------------------------------------------------------
+@pytest.mark.parametrize(
+    ("database_name", "identifier", "expected_identifier", "url"),
+    [
+        (
+            ExternalDatabaseName.PDB,
+            "1K79",
+            "1K79",
+            "https://www.rcsb.org/structure/1K79",
+        ),
+        (ExternalDatabaseName.PDB, 1234, "1234", None),
+        (
+            ExternalDatabaseName.UNIPROT,
+            "P06213",
+            "P06213",
+            "https://www.uniprot.org/uniprotkb/P06213/entry",
+        ),
+        (ExternalDatabaseName.UNIPROT, 123456, "123456", None),
+    ],
+)
+def test_external_identifier_creation(
+    database_name, identifier, expected_identifier, url
+):
+    """Test creation of ExternalIdentifier instances."""
+    external_identifier = ExternalIdentifier(
+        database_name=database_name,
+        identifier=identifier,
+        url=url,
+    )
+    assert external_identifier.database_name == database_name
+    assert external_identifier.identifier == expected_identifier
+    assert external_identifier.url == url
+
+
+def test_invalid_database_name_in_external_identifiers():
+    """Test invalid database names."""
+    # Invalid database name
+    with pytest.raises(ValidationError):
+        ExternalIdentifier(
+            database_name="INVALID_DB",  # type: ignore
+            identifier="1ABC",
+        )
+    with pytest.raises(AttributeError):
+        ExternalIdentifier(
+            database_name=ExternalDatabaseName.DUMMY,  # type: ignore
+            identifier="1ABC",
+        )