update doc strings, refine prop key / nullable relationship

examples/end_to_end/multi_file/countries/countries.ipynb

@@ -9,7 +9,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {

neo4j_runway/utils/data/data_dictionary/column.py

@@ -60,9 +60,11 @@ def validate_aliases(
 
     @model_validator(mode="after")
     def validate_model(self) -> "Column":
-        if self.primary_key or self.foreign_key:
+        if self.primary_key:
             self.nullable = False
             self.key = True
+        elif self.foreign_key:
+            self.key = True
         if self.primary_key and self.foreign_key:
             raise ValueError("`primary_key` and `foreign_key` may not both be True.")
         return self

neo4j_runway/utils/data/data_dictionary/data_dictionary.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List
 
 from pydantic import BaseModel
 
@@ -7,6 +7,15 @@
 
 
 class DataDictionary(BaseModel):
+    """
+    The data dictionary describing all tables in the data.
+
+    Attributes
+    ----------
+    table_schemas : List[TableSchema]
+        A list of TableSchema objects.
+    """
+
     table_schemas: List[TableSchema]
 
     @property

neo4j_runway/utils/data/data_dictionary/table_schema.py

@@ -99,6 +99,23 @@ def compact_dict(self) -> Dict[str, Any]:
         [compact_info.update(c.compact_dict) for c in self.columns]
         return {self.name: compact_info}
 
+    def get_column(self, column_name: str) -> Optional[Column]:
+        """
+        Retrieve a `Column` object by name.
+
+        Parameters
+        ----------
+        column_name : str
+            The name of the column.
+
+        Returns
+        -------
+        Optional[Column]
+            A `Column` object if it exists, or None
+        """
+
+        return self.columns_dict.get(column_name)
+
     def get_description(self, column_name: str) -> str:
         """
         Retrieve the description for a column. If there is no description, then return an empty string.

tests/unit/utils/data/data_dictionary/test_column.py

@@ -33,7 +33,7 @@ def test_init_nullable_and_primary_key() -> None:
 def test_init_nullable_and_foreign_key() -> None:
     c = Column(name="col_a", foreign_key=True, nullable=True)
 
-    assert not c.nullable
+    assert c.nullable
     assert c.foreign_key
 
 

tests/unit/utils/data/data_dictionary/test_table_schema.py

@@ -97,3 +97,21 @@ def test_compact_dict() -> None:
     assert len(ts.compact_dict.get(ts.name)) == 2
     assert ts.name in ts.compact_dict.keys()
     assert len(ts.compact_dict.keys()) == 1
+
+
+def test_get_column_exists() -> None:
+    ts = TableSchema(
+        name="table_a",
+        columns=[Column(name="col_a", primary_key=True), Column(name="col_b")],
+    )
+
+    assert ts.get_column("col_a") is not None
+
+
+def test_get_column_not_exists() -> None:
+    ts = TableSchema(
+        name="table_a",
+        columns=[Column(name="col_a", primary_key=True), Column(name="col_b")],
+    )
+
+    assert ts.get_column("error") is None