Allow IDGenerator be used for Primary and Alternate Keys

sdv-dev · Apr 4, 2024 · 0df558a · 0df558a
1 parent 70d6c8b
commit 0df558a
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/sdv/data_processing/data_processor.py b/sdv/data_processing/data_processor.py
@@ -644,11 +644,12 @@ def update_transformers(self, column_name_to_transformer):
             )
 
         for column, transformer in column_name_to_transformer.items():
-            if column in self._keys and not type(transformer) in (AnonymizedFaker, RegexGenerator):
+            if column in self._keys and not type(transformer) \
+                    in (AnonymizedFaker, RegexGenerator, IDGenerator):
                 raise SynthesizerInputError(
                     f"Invalid transformer '{transformer.__class__.__name__}' for a primary "
-                    f"or alternate key '{column}'. Please use 'AnonymizedFaker' or "
-                    "'RegexGenerator' instead."
+                    f"or alternate key '{column}'. Please use 'AnonymizedFaker', "
+                    "'IDGenerator', or 'RegexGenerator' instead."
                 )
 
         with warnings.catch_warnings():

diff --git a/tests/unit/data_processing/test_data_processor.py b/tests/unit/data_processing/test_data_processor.py
@@ -1588,7 +1588,10 @@ def test_update_transformers_ignores_rdt_refit_warning(self):
             dp.update_transformers({'col1': GaussianNormalizer()})
 
     def test_update_transformers_for_key(self):
-        """Test when ``transformer`` is not ``AnonymizedFaker`` or ``RegexGenerator`` for keys."""
+        """
+        Test when ``transformer`` is not ``AnonymizedFaker``, ``IDGenerator,
+        or ``RegexGenerator`` for keys.
+        """
         # Setup
         dp = DataProcessor(SingleTableMetadata())
         dp._keys = ['pk_column', 'b']
@@ -1597,7 +1600,7 @@ def test_update_transformers_for_key(self):
         # Run and Assert
         error_msg = re.escape(
             "Invalid transformer 'FloatFormatter' for a primary or alternate key 'pk_column'. "
-            "Please use 'AnonymizedFaker' or 'RegexGenerator' instead."
+            "Please use 'AnonymizedFaker', 'IDGenerator', or 'RegexGenerator' instead."
         )
         with pytest.raises(SynthesizerInputError, match=error_msg):
             dp.update_transformers({'pk_column': FloatFormatter()})