Function to exclude edges with input-row in table.

Author: MrCurtis

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 [project]
 name = "fk-graph"
 packages = ["src/fk_graph"]
-version = "0.0.11"
+version = "0.0.12"
 authors = [
   { name="Andrew Curtis", email="fk.graph@fastmail.com" },
   { name="John C Thomas" },

src/fk_graph/edge_excluders.py

@@ -0,0 +1,34 @@
+"""Edge Excluders.
+
+Functions to pass in as an `exclude_edge` argument to `get_graph`.
+
+The functions in this module that do not take any arguments can be passed in
+directly to `get_graph`. Those that do take arguments should be called with the
+required arguments, and the returned function passed to `get_graph`.
+"""
+
+def input_row_is_in_tables(table_names):
+    """Checks whether input row-node is in list of tables.
+
+    This is typically used in cases where one wishes to include rows from a
+    table as nodes in the graph, but *not* include the subsequent rows which
+    can only be reached via these nodes.
+    For example, if there is a `user` table with f-k relation to a `country`
+    table, one might want to include the country that a user is from, but not
+    also include all the other users that are from that country.
+
+    Args:
+      table_names: A list of table names.
+
+    Returns:
+      A function of the form (input_row, output_row) -> bool. The function will
+      return `True` if and only if `input_row` is from a table with one of the
+      names in `table_name`.
+    """
+    def f(input_row, output_row):
+        return _get_table_name_from_row(input_row) in table_names
+
+    return f
+
+def _get_table_name_from_row(row):
+    return row.__table__.name

tests/test_edge_excluders.py

@@ -0,0 +1,59 @@
+from unittest import TestCase
+
+from sqlalchemy import create_engine
+from sqlalchemy.orm import (
+    DeclarativeBase,
+    Mapped,
+    mapped_column,
+    Session,
+)
+
+from fk_graph.edge_excluders import input_row_is_in_tables
+
+class TestInputRowIsInTable(TestCase):
+
+    def setUp(self):
+        class Base(DeclarativeBase):
+            pass
+
+        class TableA(Base):
+            __tablename__ = "table_a"
+            id: Mapped[int] = mapped_column(primary_key=True)
+
+        class TableB(Base):
+            __tablename__ = "table_b"
+            id: Mapped[int] = mapped_column(primary_key=True)
+
+        self.engine = create_engine("sqlite+pysqlite:///:memory:")
+        Base.metadata.create_all(self.engine)
+        self.TableA, self.TableB = TableA, TableB
+
+    def test_function_returns_true_if_input_row_from_table(self):
+        with Session(self.engine) as session:
+            table_a_row = self.TableA(id=1)
+            table_b_row = self.TableB(id=1)
+            session.add_all([table_a_row, table_b_row])
+            session.commit()
+
+        with Session(self.engine) as session:
+            table_a_row = session.get_one(self.TableA, 1)
+            table_b_row = session.get_one(self.TableB, 1)
+
+            self.assertTrue(
+                input_row_is_in_tables(["table_a"])(table_a_row, table_b_row)
+            )
+
+    def test_function_returns_false_if_input_row_not_from_table(self):
+        with Session(self.engine) as session:
+            table_a_row = self.TableA(id=1)
+            table_b_row = self.TableB(id=1)
+            session.add_all([table_a_row, table_b_row])
+            session.commit()
+
+        with Session(self.engine) as session:
+            table_a_row = session.get_one(self.TableA, 1)
+            table_b_row = session.get_one(self.TableB, 1)
+
+            self.assertFalse(
+                input_row_is_in_tables(["table_b"])(table_a_row, table_b_row)
+            )