From 8a1aca92fc957e49d84839ffa58e73a516b47829 Mon Sep 17 00:00:00 2001 From: Plamen Valentinov Kolev Date: Wed, 11 Mar 2026 17:49:19 +0100 Subject: [PATCH 1/3] Update docstring --- pyproject.toml | 1 + sdv/metadata/metadata.py | 9 ++++++++- tests/unit/metadata/test_metadata.py | 19 +++++++++++++++++++ 3 files changed, 28 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 4cf85e467..adde43458 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -202,6 +202,7 @@ exclude = [ "*.ipynb", ".ipynb_checkpoints", "tasks.py", + "*.md", ] [tool.ruff.lint] diff --git a/sdv/metadata/metadata.py b/sdv/metadata/metadata.py index abbda2b68..1da24670d 100644 --- a/sdv/metadata/metadata.py +++ b/sdv/metadata/metadata.py @@ -432,7 +432,14 @@ def add_column_relationship( super().add_column_relationship(table_name, relationship_type, column_names) def set_primary_key(self, column_name, table_name=None): - """Set the primary key of a table.""" + """Set the primary key of a table. + + Args: + table_name (str): + Name of the table to set the primary key. + column_name (str, tulple[str]): + Name (or tuple of names) of the primary key column(s). + """ table_name = self._handle_table_name(table_name) super().set_primary_key(table_name, column_name) diff --git a/tests/unit/metadata/test_metadata.py b/tests/unit/metadata/test_metadata.py index d6025a4ec..aa09d9818 100644 --- a/tests/unit/metadata/test_metadata.py +++ b/tests/unit/metadata/test_metadata.py @@ -1319,3 +1319,22 @@ def test_remove_column_alternate_key(self): assert list(trades_mock.columns.keys()) == ['id', 'cost', 'quantity', 'time'] assert trades_mock.alternate_keys == [] assert metadata._multi_table_updated + + def test_set_primary_key_docstring(self): + """Test that the `set_primary_key` docstring contains expected info.""" + # Setup + expected_msg = ( + 'Set the primary key of a table.\n\n' + 'Args:\n' + ' table_name (str):\n' + ' Name of the table to set the primary key.\n' + ' column_name (str, tulple[str]):\n' + ' Name (or tuple of names) of the primary key column(s).\n' + ) + metadata = Metadata() + + # Run + result = metadata.set_primary_key.__doc__ + + # Assert + assert result == expected_msg From 682dc82ccaf2baba10488f9e7020ba349370121b Mon Sep 17 00:00:00 2001 From: Plamen Valentinov Kolev Date: Thu, 12 Mar 2026 16:15:13 +0100 Subject: [PATCH 2/3] Update docstrings --- pyproject.toml | 1 - sdv/metadata/metadata.py | 4 ++-- sdv/metadata/multi_table.py | 4 ++-- sdv/metadata/single_table.py | 4 ++-- tests/unit/metadata/test_metadata.py | 4 ++-- 5 files changed, 8 insertions(+), 9 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index adde43458..4cf85e467 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -202,7 +202,6 @@ exclude = [ "*.ipynb", ".ipynb_checkpoints", "tasks.py", - "*.md", ] [tool.ruff.lint] diff --git a/sdv/metadata/metadata.py b/sdv/metadata/metadata.py index 1da24670d..ebea2ea55 100644 --- a/sdv/metadata/metadata.py +++ b/sdv/metadata/metadata.py @@ -435,10 +435,10 @@ def set_primary_key(self, column_name, table_name=None): """Set the primary key of a table. Args: + column_name (str, list[str]): + Name (or list of names) of the primary key column(s). table_name (str): Name of the table to set the primary key. - column_name (str, tulple[str]): - Name (or tuple of names) of the primary key column(s). """ table_name = self._handle_table_name(table_name) super().set_primary_key(table_name, column_name) diff --git a/sdv/metadata/multi_table.py b/sdv/metadata/multi_table.py index 6251a458c..39f1d34c2 100644 --- a/sdv/metadata/multi_table.py +++ b/sdv/metadata/multi_table.py @@ -689,8 +689,8 @@ def set_primary_key(self, table_name, column_name): Args: table_name (str): Name of the table to set the primary key. - column_name (str, tulple[str]): - Name (or tuple of names) of the primary key column(s). + column_name (str, list[str]): + Name (or list of names) of the primary key column(s). """ self._validate_table_exists(table_name) self.tables[table_name].set_primary_key(column_name) diff --git a/sdv/metadata/single_table.py b/sdv/metadata/single_table.py index 91a336e17..c0bb49bfa 100644 --- a/sdv/metadata/single_table.py +++ b/sdv/metadata/single_table.py @@ -855,8 +855,8 @@ def set_primary_key(self, column_name): """Set the metadata primary key. Args: - column_name (str): - Name of the primary key column(s). + column_name (str, list[str]): + Name (or list of names) of the primary key column(s). """ if isinstance(column_name, list) and len(column_name) == 1: column_name = column_name[0] diff --git a/tests/unit/metadata/test_metadata.py b/tests/unit/metadata/test_metadata.py index aa09d9818..6ce314560 100644 --- a/tests/unit/metadata/test_metadata.py +++ b/tests/unit/metadata/test_metadata.py @@ -1326,10 +1326,10 @@ def test_set_primary_key_docstring(self): expected_msg = ( 'Set the primary key of a table.\n\n' 'Args:\n' + ' column_name (str, list[str]):\n' + ' Name (or list of names) of the primary key column(s).\n' ' table_name (str):\n' ' Name of the table to set the primary key.\n' - ' column_name (str, tulple[str]):\n' - ' Name (or tuple of names) of the primary key column(s).\n' ) metadata = Metadata() From 61337c04d938cc6cb8127f8326a7adf65ec32aaa Mon Sep 17 00:00:00 2001 From: Plamen Valentinov Kolev Date: Thu, 12 Mar 2026 17:09:02 +0100 Subject: [PATCH 3/3] Remove flaky test --- tests/unit/metadata/test_metadata.py | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/tests/unit/metadata/test_metadata.py b/tests/unit/metadata/test_metadata.py index 6ce314560..d6025a4ec 100644 --- a/tests/unit/metadata/test_metadata.py +++ b/tests/unit/metadata/test_metadata.py @@ -1319,22 +1319,3 @@ def test_remove_column_alternate_key(self): assert list(trades_mock.columns.keys()) == ['id', 'cost', 'quantity', 'time'] assert trades_mock.alternate_keys == [] assert metadata._multi_table_updated - - def test_set_primary_key_docstring(self): - """Test that the `set_primary_key` docstring contains expected info.""" - # Setup - expected_msg = ( - 'Set the primary key of a table.\n\n' - 'Args:\n' - ' column_name (str, list[str]):\n' - ' Name (or list of names) of the primary key column(s).\n' - ' table_name (str):\n' - ' Name of the table to set the primary key.\n' - ) - metadata = Metadata() - - # Run - result = metadata.set_primary_key.__doc__ - - # Assert - assert result == expected_msg