Create MultiTableMetadata class (#924)

* Create MultiTableMetadata class.

* Increase coverage to 100%

* Fix typo

sdv/metadata/__init__.py

@@ -3,13 +3,15 @@
 from sdv.metadata import visualization
 from sdv.metadata.dataset import Metadata
 from sdv.metadata.errors import MetadataError, MetadataNotFittedError
+from sdv.metadata.multi_table import MultiTableMetadata
 from sdv.metadata.single_table import SingleTableMetadata
 from sdv.metadata.table import Table
 
 __all__ = (
     'Metadata',
     'MetadataError',
     'MetadataNotFittedError',
+    'MultiTableMetadata',
     'SingleTableMetadata',
     'Table',
     'visualization'

sdv/metadata/multi_table.py

@@ -0,0 +1,56 @@
+"""Multi Table Metadata."""
+
+import json
+from copy import deepcopy
+
+from sdv.metadata.single_table import SingleTableMetadata
+
+
+class MultiTableMetadata:
+    """Multi Table Metadata class."""
+
+    def __init__(self):
+        self._tables = {}
+        self._relationships = []
+
+    def to_dict(self):
+        """Return a python ``dict`` representation of the ``MultiTableMetadata``."""
+        metadata = {'tables': {}, 'relationships': []}
+        for table_name, single_table_metadata in self._tables.items():
+            metadata['tables'][table_name] = single_table_metadata.to_dict()
+
+        metadata['relationships'] = deepcopy(self._relationships)
+        return metadata
+
+    def _set_metadata_dict(self, metadata):
+        """Set a ``metadata`` dictionary to the current instance.
+
+        Args:
+            metadata (dict):
+                Python dictionary representing a ``MultiTableMetadata`` object.
+        """
+        for table_name, table_dict in metadata.get('tables', {}).items():
+            self._tables[table_name] = SingleTableMetadata._load_from_dict(table_dict)
+
+        for relationship in metadata.get('relationships', []):
+            self._relationships.append(relationship)
+
+    @classmethod
+    def _load_from_dict(cls, metadata):
+        """Create a ``MultiTableMetadata`` instance from a python ``dict``.
+
+        Args:
+            metadata (dict):
+                Python dictionary representing a ``MultiTableMetadata`` object.
+
+        Returns:
+            Instance of ``MultiTableMetadata``.
+        """
+        instance = cls()
+        instance._set_metadata_dict(metadata)
+        return instance
+
+    def __repr__(self):
+        """Pretty print the ``MultiTableMetadata``."""
+        printed = json.dumps(self.to_dict(), indent=4)
+        return printed

tests/integration/metadata/test_multi_table.py

@@ -0,0 +1,21 @@
+"""Integration tests for Multi Table Metadata."""
+
+from sdv.metadata import MultiTableMetadata
+
+
+def test_single_table_metadata():
+    """Test ``MultiTableMetadata``."""
+
+    # Create an instance
+    instance = MultiTableMetadata()
+
+    # To dict
+    result = instance.to_dict()
+
+    # Assert
+    assert result == {
+        'tables': {},
+        'relationships': []
+    }
+    assert instance._tables == {}
+    assert instance._relationships == []

tests/unit/metadata/test_multi_table.py

@@ -0,0 +1,237 @@
+"""Test Multi Table Metadata."""
+
+from unittest.mock import Mock, patch
+
+from sdv.metadata.multi_table import MultiTableMetadata
+
+
+class TestMultiTableMetadata:
+    """Test ``MultiTableMetadata`` class."""
+
+    def test___init__(self):
+        """Test the ``__init__`` method of ``MultiTableMetadata``."""
+        # Run
+        instance = MultiTableMetadata()
+
+        # Assert
+        assert instance._tables == {}
+        assert instance._relationships == []
+
+    def test_to_dict(self):
+        """Test the ``to_dict`` method of ``MultiTableMetadata``.
+
+        Setup:
+            - Instance of ``MultiTableMetadata``.
+            - Add mocked values to ``instance._tables`` and ``instance._relationships``.
+        Mock:
+            - Mock ``SingleTableMetadata`` like object to ``instance._tables``.
+
+        Output:
+            - A dict representation containing ``tables`` and ``relationships`` has to be returned
+              with ``dict`` values of ``tables``.
+        """
+        # Setup
+        table_accounts = Mock()
+        table_accounts.to_dict.return_value = {
+            'id': {'sdtype': 'numerical'},
+            'branch_id': {'sdtype': 'numerical'},
+            'amount': {'sdtype': 'numerical'},
+            'start_date': {'sdtype': 'datetime'},
+            'owner': {'sdtype': 'text'},
+        }
+        table_branches = Mock()
+        table_branches.to_dict.return_value = {
+            'id': {'sdtype': 'numerical'},
+            'name': {'sdtype': 'text'},
+        }
+        instance = MultiTableMetadata()
+        instance._tables = {
+            'accounts': table_accounts,
+            'branches': table_branches
+        }
+        instance._relationships = [{
+            'parent_table_name': 'accounts',
+            'parent_primary_key': 'id',
+            'child_table_name': 'branches',
+            'chil_foreign_key': 'branch_id',
+        }]
+
+        # Run
+        result = instance.to_dict()
+
+        # Assert
+        expected_result = {
+            'tables': {
+                'accounts': {
+                    'id': {'sdtype': 'numerical'},
+                    'branch_id': {'sdtype': 'numerical'},
+                    'amount': {'sdtype': 'numerical'},
+                    'start_date': {'sdtype': 'datetime'},
+                    'owner': {'sdtype': 'text'},
+                },
+                'branches': {
+                    'id': {'sdtype': 'numerical'},
+                    'name': {'sdtype': 'text'},
+                }
+            },
+            'relationships': [{
+                'parent_table_name': 'accounts',
+                'parent_primary_key': 'id',
+                'child_table_name': 'branches',
+                'chil_foreign_key': 'branch_id',
+            }]
+        }
+        assert result == expected_result
+
+    @patch('sdv.metadata.multi_table.SingleTableMetadata')
+    def test__set_metadata(self, mock_singletablemetadata):
+        """Test the ``_set_metadata`` method for ``MultiTableMetadata``.
+
+        Setup:
+            - instance of ``MultiTableMetadata``.
+            - A dict representing a ``MultiTableMetadata``.
+
+        Mock:
+            - Mock ``SingleTableMetadata`` from ``sdv.metadata.multi_table``
+
+        Side Effects:
+            - ``instance`` now contains ``instance._tables`` and ``instance._relationships``.
+            - ``SingleTableMetadata._load_from_dict`` has been called.
+        """
+        # Setup
+        multitable_metadata = {
+            'tables': {
+                'accounts': {
+                    'id': {'sdtype': 'numerical'},
+                    'branch_id': {'sdtype': 'numerical'},
+                    'amount': {'sdtype': 'numerical'},
+                    'start_date': {'sdtype': 'datetime'},
+                    'owner': {'sdtype': 'text'},
+                },
+                'branches': {
+                    'id': {'sdtype': 'numerical'},
+                    'name': {'sdtype': 'text'},
+                }
+            },
+            'relationships': [{
+                'parent_table_name': 'accounts',
+                'parent_primary_key': 'id',
+                'child_table_name': 'branches',
+                'chil_foreign_key': 'branch_id',
+            }]
+        }
+
+        single_table_accounts = object()
+        single_table_branches = object()
+        mock_singletablemetadata._load_from_dict.side_effect = [
+            single_table_accounts,
+            single_table_branches
+        ]
+
+        instance = MultiTableMetadata()
+
+        # Run
+        instance._set_metadata_dict(multitable_metadata)
+
+        # Assert
+        assert instance._tables == {
+            'accounts': single_table_accounts,
+            'branches': single_table_branches
+        }
+
+        assert instance._relationships == [{
+            'parent_table_name': 'accounts',
+            'parent_primary_key': 'id',
+            'child_table_name': 'branches',
+            'chil_foreign_key': 'branch_id',
+        }]
+
+    @patch('sdv.metadata.multi_table.SingleTableMetadata')
+    def test__load_from_dict(self, mock_singletablemetadata):
+        """Test that ``_load_from_dict`` returns a instance of ``MultiTableMetadata``.
+
+        Test that when calling the ``_load_from_dict`` method a new instance with the passed
+        python ``dict`` details should be created.
+
+        Setup:
+            - A dict representing a ``MultiTableMetadata``.
+
+        Mock:
+            - Mock ``SingleTableMetadata`` from ``sdv.metadata.multi_table``
+
+        Output:
+            - ``instance`` that contains ``instance._tables`` and ``instance._relationships``.
+
+        Side Effects:
+            - ``SingleTableMetadata._load_from_dict`` has been called.
+        """
+        # Setup
+        multitable_metadata = {
+            'tables': {
+                'accounts': {
+                    'id': {'sdtype': 'numerical'},
+                    'branch_id': {'sdtype': 'numerical'},
+                    'amount': {'sdtype': 'numerical'},
+                    'start_date': {'sdtype': 'datetime'},
+                    'owner': {'sdtype': 'text'},
+                },
+                'branches': {
+                    'id': {'sdtype': 'numerical'},
+                    'name': {'sdtype': 'text'},
+                }
+            },
+            'relationships': [{
+                'parent_table_name': 'accounts',
+                'parent_primary_key': 'id',
+                'child_table_name': 'branches',
+                'chil_foreign_key': 'branch_id',
+            }]
+        }
+
+        single_table_accounts = object()
+        single_table_branches = object()
+        mock_singletablemetadata._load_from_dict.side_effect = [
+            single_table_accounts,
+            single_table_branches
+        ]
+
+        # Run
+        instance = MultiTableMetadata._load_from_dict(multitable_metadata)
+
+        # Assert
+        assert instance._tables == {
+            'accounts': single_table_accounts,
+            'branches': single_table_branches
+        }
+
+        assert instance._relationships == [{
+            'parent_table_name': 'accounts',
+            'parent_primary_key': 'id',
+            'child_table_name': 'branches',
+            'chil_foreign_key': 'branch_id',
+        }]
+
+    @patch('sdv.metadata.multi_table.json')
+    def test___repr__(self, mock_json):
+        """Test that the ``__repr__`` method.
+
+        Test that the ``__repr__`` method calls the ``json.dumps``  method and
+        returns its output.
+
+        Setup:
+            - Instance of ``MultiTableMetadata``.
+        Mock:
+            - ``json`` from ``sdv.metadata.multi_table``.
+
+        Output:
+            - ``json.dumps`` return value.
+        """
+        # Setup
+        instance = MultiTableMetadata()
+
+        # Run
+        res = instance.__repr__()
+
+        # Assert
+        mock_json.dumps.assert_called_once_with(instance.to_dict(), indent=4)
+        assert res == mock_json.dumps.return_value