cleanup

Author: gsheni

sdv/data_processing/data_processor.py

@@ -780,8 +780,6 @@ def update_transformers(self, column_name_to_transformer):
             )
 
         for column, transformer in column_name_to_transformer.items():
-            if transformer is None:
-                continue
             if column in self._keys and not transformer.is_generator():
                 raise SynthesizerInputError(
                     f"Invalid transformer '{transformer.__class__.__name__}' for a primary "

sdv/multi_table/base.py

@@ -474,14 +474,12 @@ def _assign_table_transformers(self, synthesizer, table_name, table_data):
         """
         synthesizer.auto_assign_transformers(table_data)
         column_name_to_transformers = {}
+
         primary_key = self.metadata.tables[table_name].primary_key
-        for relation in self.metadata.relationships:
-            column_name = relation['child_foreign_key']
-            if (
-                relation['child_table_name'] == table_name
-                and relation['child_foreign_key'] != primary_key
-            ):
-                column_name_to_transformers[column_name] = None
+        foreign_keys = self.metadata._get_all_foreign_keys(table_name)
+        for foreign_key in foreign_keys:
+            if foreign_key != primary_key:
+                column_name_to_transformers[foreign_key] = None
         synthesizer.update_transformers(column_name_to_transformers)
 
     def auto_assign_transformers(self, data):

tests/unit/multi_table/test_base.py

@@ -689,6 +689,26 @@ def test_auto_assign_transformers_foreign_key_none(self):
         nesreca_synthesizer.update_transformers.assert_called_once_with({'a': None, 'b': None})
         oseba_synthesizer.update_transformers.assert_called_once_with({'a': None, 'b': None})
 
+    def test_auto_assign_transformers_primary_key_not_assigned(self):
+        """Test that if table has a PK (which also an FK), it is not assigned a Transformer."""
+        # Setup
+        metadata = get_multi_table_metadata()
+        instance = BaseMultiTableSynthesizer(metadata)
+        data = {'nesreca': Mock(), 'oseba': Mock()}
+        instance.metadata.tables['nesreca'].primary_key = 'a'
+        instance.metadata._get_all_foreign_keys = Mock(return_value=['a', 'b'])
+        nesreca_synthesizer = Mock()
+        oseba_synthesizer = Mock()
+        instance._table_synthesizers['nesreca'] = nesreca_synthesizer
+        instance._table_synthesizers['oseba'] = oseba_synthesizer
+
+        # Run
+        instance.auto_assign_transformers(data)
+
+        # Assert
+        nesreca_synthesizer.update_transformers.assert_called_once_with({'b': None})
+        oseba_synthesizer.update_transformers.assert_called_once_with({'a': None, 'b': None})
+
     def test_auto_assign_transformers_missing_table(self):
         """Test it errors out when the passed table was not seen in the metadata."""
         # Setup