From c519913f10cf9d36c991496d14095059b953f9d9 Mon Sep 17 00:00:00 2001
From: David Grayson <dgrayson@absci.com>
Date: Mon, 6 Jan 2025 23:52:13 +0000
Subject: [PATCH] propagate allow_missing_coord and add test

---
 src/biotite/structure/info/atoms.py      | 13 ++++++++--
 src/biotite/structure/io/pdbx/convert.py | 22 ++++++++---------
 tests/structure/test_info.py             | 30 ++++++++++++++++++++++++
 3 files changed, 52 insertions(+), 13 deletions(-)

diff --git a/src/biotite/structure/info/atoms.py b/src/biotite/structure/info/atoms.py
index a404b39bc..3e78b776d 100644
--- a/src/biotite/structure/info/atoms.py
+++ b/src/biotite/structure/info/atoms.py
@@ -18,7 +18,7 @@
 # fmt: on
 
 
-def residue(res_name):
+def residue(res_name, allow_missing_coord=False):
     """
     Get an atom array, representing the residue with the given name.
 
@@ -30,6 +30,11 @@ def residue(res_name):
     ----------
     res_name : str
         The up to 3-letter name of the residue.
+    allow_missing_coord: bool, optional
+        Whether to allow missing coordinate values in the residue.
+        If ``True``, these will be represented as ``nan`` values.
+        If ``False``, a ``ValueError`` is raised when missing coordinates
+        are encountered.
 
     Returns
     -------
@@ -74,7 +79,11 @@ def residue(res_name):
     from biotite.structure.io.pdbx import get_component
 
     try:
-        component = get_component(get_ccd(), res_name=res_name)
+        component = get_component(
+            get_ccd(),
+            res_name=res_name,
+            allow_missing_coord=allow_missing_coord,
+        )
     except KeyError:
         raise KeyError(f"No atom information found for residue '{res_name}' in CCD")
     component.hetero[:] = res_name not in NON_HETERO_RESIDUES
diff --git a/src/biotite/structure/io/pdbx/convert.py b/src/biotite/structure/io/pdbx/convert.py
index 2d421ae28..7a1078104 100644
--- a/src/biotite/structure/io/pdbx/convert.py
+++ b/src/biotite/structure/io/pdbx/convert.py
@@ -1186,7 +1186,7 @@ def get_component(
     data_block=None,
     use_ideal_coord=True,
     res_name=None,
-    allow_missing_coords=False,
+    allow_missing_coord=False,
 ):
     """
     Create an :class:`AtomArray` for a chemical component from the
@@ -1215,10 +1215,10 @@ def get_component(
         In this case, the component with the given residue name is
         read.
         By default, all rows would be read in this case.
-    allow_missing_coords: bool
+    allow_missing_coord: bool, optional
         Whether to allow missing coordinate values in components.
-        If `True`, these will be represented as `nan` values.
-        If `False`, a `ValueError` is raised when missing coordinates
+        If ``True``, these will be represented as ``nan`` values.
+        If ``False``, a ``ValueError`` is raised when missing coordinates
         are encountered.
 
     Returns
@@ -1311,7 +1311,7 @@ def get_component(
             raise
         array.coord = _parse_component_coordinates(
             [atom_category[field] for field in alt_coord_fields],
-            keep_missing=allow_missing_coords,
+            allow_missing=allow_missing_coord,
         )
 
     try:
@@ -1342,19 +1342,19 @@ def get_component(
     return array
 
 
-def _parse_component_coordinates(coord_columns, keep_missing=False):
+def _parse_component_coordinates(coord_columns, allow_missing=False):
     coord = np.zeros((len(coord_columns[0]), 3), dtype=np.float32)
     for i, column in enumerate(coord_columns):
         if column.mask is not None and column.mask.array.any():
-            if not keep_missing:
-                raise ValueError(
-                    "Missing coordinates for some atoms",
-                )
-            else:
+            if allow_missing:
                 warnings.warn(
                     "Missing coordinates for some atoms. Those will be set to nan",
                     UserWarning,
                 )
+            else:
+                raise ValueError(
+                    "Missing coordinates for some atoms",
+                )
         coord[:, i] = column.as_array(np.float32, masked_value=np.nan)
     return coord
 
diff --git a/tests/structure/test_info.py b/tests/structure/test_info.py
index 90b9cbc90..99068f60a 100644
--- a/tests/structure/test_info.py
+++ b/tests/structure/test_info.py
@@ -192,3 +192,33 @@ def test_set_ccd_path(fake_ccd_path):
 
     # The new fake CCD has only a single compound
     assert strucinfo.all_residues() == ["FOO"]
+
+
+@pytest.mark.parametrize(
+    "res_name, allow_missing_coord, should_raise",
+    [
+        ("ALA", False, False),
+        ("ALA", True, False),
+        ("A1IQW", True, False),
+        ("A1IQW", False, True),
+        ("RRE", True, False),
+        ("RRE", False, True),
+    ],
+)
+def test_residue(res_name, allow_missing_coord, should_raise):
+    """
+    Test if the residue function returns an atom array or not.
+    ALA --> standard amino acid, yes in both conditions
+    A1IQW --> yes only with allow_missing_coord=True
+    RRE --> yes only with allow_missing_coord=True
+    Make sure correct exceptions are raised when the non-standard residue
+    is used with allow_missing_coord=False.
+    """
+    if should_raise:
+        with pytest.raises(ValueError):
+            strucinfo.residue(res_name, allow_missing_coord=allow_missing_coord)
+    else:
+        result = strucinfo.residue(res_name, allow_missing_coord=allow_missing_coord)
+        assert isinstance(result, struc.AtomArray)
+        assert result.array_length() > 0
+        assert np.all(result.res_name == res_name)