Make extra field check optional and improve logic

monarch-initiative · Nov 11, 2024 · 9b376a2 · 9b376a2
1 parent 02da771
commit 9b376a2
Show file tree

Hide file tree

Showing 5 changed files with 9 additions and 9 deletions.
diff --git a/src/koza/io/writer/writer.py b/src/koza/io/writer/writer.py
@@ -19,7 +19,7 @@ def __init__(
         node_properties: List[str] = None,
         edge_properties: List[str] = None,
         sssom_config: SSSOMConfig = None,
-        skip_checks: bool = False,
+        check_fields: bool = False,
         kwargs: dict = None,
     ):
         """Do not override this method; implement `init` instead."""
@@ -28,7 +28,7 @@ def __init__(
         self.node_properties = node_properties
         self.edge_properties = edge_properties
         self.sssom_config = sssom_config
-        self.skip_checks = skip_checks
+        self.check_fields = check_fields
         self.converter = KGXConverter()
 
         kwargs = kwargs or {}
@@ -39,15 +39,15 @@ def write(self, entities: Iterable):
 
         if nodes:
             for node in nodes:
-                if not self.skip_checks:
+                if self.check_fields:
                     self.check_extra_fields(tuple(node.keys()), tuple(self.node_properties))
                 self.write_node(node)
 
         if edges:
             for edge in edges:
                 if self.sssom_config:
                     edge = self.sssom_config.apply_mapping(edge)
-                if not self.skip_checks:
+                if self.check_fields:
                     self.check_extra_fields(tuple(edge.keys()), tuple(self.edge_properties))
                 self.write_edge(edge)
 

diff --git a/tests/unit/test_tsvwriter_node_and_edge.py b/tests/unit/test_tsvwriter_node_and_edge.py
@@ -97,7 +97,7 @@ def test_tsv_writer():
     outdir = "output/tests"
     outfile = "tsvwriter-node-and-edge"
 
-    t = TSVWriter(outdir, outfile, node_properties, edge_properties)
+    t = TSVWriter(outdir, outfile, node_properties, edge_properties, check_fields=True)
     t.write(ent)
     t.finalize()
 

diff --git a/tests/unit/test_tsvwriter_node_and_edge_extra_params.py b/tests/unit/test_tsvwriter_node_and_edge_extra_params.py
@@ -89,7 +89,7 @@ def test_tsv_writer_extra_node_params():
     outdir = "output/tests"
     outfile = "tsvwriter-node-and-edge"
 
-    t = TSVWriter(outdir, outfile, node_properties, edge_properties)
+    t = TSVWriter(outdir, outfile, node_properties, edge_properties, check_fields=True)
     expected_message = "Extra fields found in row: ['deprecated', 'has_attribute', 'name']"
     with pytest.raises(ValueError, match=re.escape(expected_message)):
         t.write(ent)
@@ -184,7 +184,7 @@ def test_tsv_writer_extra_edge_params():
     outdir = "output/tests"
     outfile = "tsvwriter-node-and-edge"
 
-    t = TSVWriter(outdir, outfile, node_properties, edge_properties)
+    t = TSVWriter(outdir, outfile, node_properties, edge_properties, check_fields=True)
     expected_message = "Extra fields found in row: ['qualifier', 'retrieval_source_ids', 'subject_aspect_qualifier']"
     with pytest.raises(ValueError, match=re.escape(expected_message)):
         t.write(ent)
diff --git a/tests/unit/test_tsvwriter_node_only.py b/tests/unit/test_tsvwriter_node_only.py
@@ -37,7 +37,7 @@ def test_tsv_writer():
     outdir = "output/tests"
     outfile = "tsvwriter-node-only"
 
-    t = TSVWriter(outdir, outfile, node_properties)
+    t = TSVWriter(outdir, outfile, node_properties, check_fields=True)
     t.write(ent)
     t.finalize()
 

diff --git a/tests/unit/test_tsvwriter_node_only_extra_params.py b/tests/unit/test_tsvwriter_node_only_extra_params.py
@@ -38,7 +38,7 @@ def test_tsv_writer():
 
     t = TSVWriter(outdir, outfile, node_properties)
 
-    t = TSVWriter(outdir, outfile, node_properties)
+    t = TSVWriter(outdir, outfile, node_properties, check_fields=True)
     expected_message = "Extra fields found in row: ['has_attribute', 'name']"
     with pytest.raises(ValueError, match=re.escape(expected_message)):
         t.write(ent)