def test_get_primary_key(self): """Test get primary key""" # Setup table_meta = {'primary_key': 'a_primary_key'} metadata = Mock(spec=Metadata) metadata.get_table_meta.return_value = table_meta # Run result = Metadata.get_primary_key(metadata, 'test') # Asserts assert result == 'a_primary_key' metadata.get_table_meta.assert_called_once_with('test')
def test_get_primary_key(self): """Test get primary key""" # Setup table_meta = { 'primary_key': 'pk' } # Run metadata = Mock() metadata.get_table_meta.return_value = table_meta table_name = 'test' result = Metadata.get_primary_key(metadata, table_name) # Asserts expected = 'pk' metadata.get_table_meta.assert_called_once_with('test') assert result == expected
class BaseRelationalModel: """Base class for all the relational models. The ``BaseRelationalModel`` class defines the common API that all the relational models need to implement, as well as common functionality. Args: metadata (dict, str or Metadata): Metadata dict, path to the metadata JSON file or Metadata instance itself. root_path (str or None): Path to the dataset directory. If ``None`` and metadata is a path, the metadata location is used. If ``None`` and metadata is a dict, the current working directory is used. """ metadata = None def __init__(self, metadata, root_path=None): if isinstance(metadata, Metadata): self.metadata = metadata else: self.metadata = Metadata(metadata, root_path) self._primary_key_generators = dict() self._remaining_primary_keys = dict() def _fit(self, tables=None): """Fit this relational model instance to the dataset data. Args: tables (dict): Dictionary with the table names as key and ``pandas.DataFrame`` instances as values. If ``None`` is given, the tables will be loaded from the paths indicated in ``metadata``. Defaults to ``None``. """ raise NotImplementedError() def fit(self, tables=None): """Fit this relational model instance to the dataset data. Args: tables (dict): Dictionary with the table names as key and ``pandas.DataFrame`` instances as values. If ``None`` is given, the tables will be loaded from the paths indicated in ``metadata``. Defaults to ``None``. """ self._fit(tables) self.fitted = True def _reset_primary_keys_generators(self): """Reset the primary key generators.""" self._primary_key_generators = dict() self._remaining_primary_keys = dict() def _get_primary_keys(self, table_name, num_rows): """Return the primary key and amount of values for the requested table. Args: table_name (str): Name of the table to get the primary keys from. num_rows (str): Number of ``primary_keys`` to generate. Returns: tuple (str, pandas.Series): primary key name and primary key values. If the table has no primary key, ``(None, None)`` is returned. Raises: ValueError: If the ``metadata`` contains invalid types or subtypes, or if there are not enough primary keys left on any of the generators. NotImplementedError: If the primary key subtype is a ``datetime``. """ primary_key = self.metadata.get_primary_key(table_name) field = self.metadata.get_fields(table_name)[primary_key] generator = self._primary_key_generators.get(table_name) if generator is None: if field['type'] != 'id': raise ValueError('Only columns with type `id` can be primary keys') subtype = field.get('subtype', 'integer') if subtype == 'integer': generator = itertools.count() remaining = np.inf elif subtype == 'string': regex = field.get('regex', r'^[a-zA-Z]+$') generator, remaining = utils.strings_from_regex(regex) elif subtype == 'datetime': raise NotImplementedError('Datetime ids are not yet supported') else: raise ValueError('Only `integer` or `string` id columns are supported.') self._primary_key_generators[table_name] = generator self._remaining_primary_keys[table_name] = remaining else: remaining = self._remaining_primary_keys[table_name] if remaining < num_rows: raise ValueError( 'Not enough unique values for primary key of table {}' ' to generate {} samples.'.format(table_name, num_rows) ) self._remaining_primary_keys[table_name] -= num_rows primary_key_values = pd.Series([x for i, x in zip(range(num_rows), generator)]) return primary_key_values def _sample(self, table_name=None, num_rows=None, sample_children=True): """Generate synthetic data for one table or the entire dataset.""" raise NotImplementedError() def sample(self, table_name=None, num_rows=None, sample_children=True, reset_primary_keys=False): """Generate synthetic data for one table or the entire dataset. If a ``table_name`` is given and ``sample_children`` is ``False``, a ``pandas.DataFrame`` with the values from the indicated table is returned. Otherwise, if ``sample_children`` is ``True``, a dictionary containing both the table and all its descendant tables is returned. If no ``table_name`` is given, the entire dataset is sampled and returned in a dictionary. If ``num_rows`` is given, the root tables of the dataset will contain the indicated number of rows. Otherwise, the number of rows will be the same as in the original dataset. Number of rows in the child tables cannot be controlled and always will depend on the values from the sampled parent tables. If ``reset_primary_keys`` is ``True``, the primary key generators will be reset. Args: table_name (str): Name of the table to sample from. If not passed, sample the entire dataset. num_rows (int): Amount of rows to sample. If ``None``, sample the same number of rows as there were in the original table. sample_children (bool): Whether or not sample child tables. Used only if ``table_name`` is given. Defaults to ``True``. reset_primary_keys (bool): Whether or not reset the primary keys generators. Defaults to ``False``. Returns: dict or pandas.DataFrame: - Returns a ``dict`` when ``sample_children`` is ``True`` with the sampled table and child tables. - Returns a ``pandas.DataFrame`` when ``sample_children`` is ``False``. Raises: NotFittedError: A ``NotFittedError`` is raised when the model has not been fitted yet. """ if not self.fitted: raise NotFittedError('SDV instance has not been fitted') if reset_primary_keys: self._reset_primary_keys_generators() return self._sample(table_name, num_rows, sample_children) def save(self, path): """Save this instance to the given path using pickle. Args: path (str): Path where the instance will be serialized. """ with open(path, 'wb') as output: pickle.dump(self, output) @classmethod def load(cls, path): """Load a model from a given path. Args: path (str): Path from which to load the instance. """ with open(path, 'rb') as f: return pickle.load(f)