sdv 1.36.3.dev0__tar.gz → 1.37.0__tar.gz

{sdv-1.36.3.dev0/sdv.egg-info → sdv-1.37.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sdv
-Version: 1.36.3.dev0
+Version: 1.37.0
 Summary: Generate synthetic data for single table, multi table and sequential data
 Author-email: "DataCebo, Inc." <info@sdv.dev>
 License-Expression: BUSL-1.1

{sdv-1.36.3.dev0 → sdv-1.37.0}/pyproject.toml

@@ -149,7 +149,7 @@ namespaces = false
 version = {attr = 'sdv.__version__'}
 
 [tool.bumpversion]
-current_version = "1.36.3.dev0"
+current_version = "1.37.0"
 parse = '(?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)(\.(?P<release>[a-z]+)(?P<candidate>\d+))?'
 serialize = [
     '{major}.{minor}.{patch}.{release}{candidate}',

{sdv-1.36.3.dev0 → sdv-1.37.0}/sdv/__init__.py

@@ -6,7 +6,7 @@
 
 __author__ = 'DataCebo, Inc.'
 __email__ = 'info@sdv.dev'
-__version__ = '1.36.3.dev0'
+__version__ = '1.37.0'
 
 
 import sys

{sdv-1.36.3.dev0 → sdv-1.37.0}/sdv/cag/_utils.py

@@ -1,3 +1,4 @@
+import importlib
 import re
 import warnings
 
@@ -227,3 +228,44 @@ def _validate_constraints_single_table(constraints, synthesizer_fitted):
             )
 
     return constraints
+
+
+def load_constraint_from_dict(constraint_dict):
+    """Load a constraint from a constraint dictionary.
+
+    Args:
+        constraint_dict (dict):
+            A constraint dictionary containing the following keys:
+            - `class_name` (str): The constraint class name.
+            - `parameters` (dict): Dictionary of the parameters used to instantiate the constraint.
+
+    Returns:
+        Instance of `class_name` constraint instantiated with the given `parameters`.
+    """
+    expected_keys = {'class_name', 'parameters'}
+    if not isinstance(constraint_dict, dict) or set(constraint_dict.keys()) != expected_keys:
+        raise ValueError(
+            'Invalid `constraint_dict`. Expected dictionary with keys `class_name` and '
+            f' `parameters`, got {constraint_dict}.'
+        )
+
+    class_name = constraint_dict['class_name']
+    parameters = constraint_dict['parameters']
+    if not isinstance(class_name, str):
+        raise ValueError('`class_name` must be a string.')
+
+    if not isinstance(parameters, dict):
+        raise ValueError('`parameters` must be a dict.')
+
+    cag_module = importlib.import_module('sdv.cag')
+    try:
+        sandbox_module = importlib.import_module('sdv.cag.sandbox')
+        sandbox_constraint = getattr(sandbox_module, class_name, None)
+    except ModuleNotFoundError:
+        sandbox_constraint = None
+
+    constraint_class = getattr(cag_module, class_name, sandbox_constraint)
+    if constraint_class is None:
+        raise ValueError(f"Unknown `constraint_class` '{class_name}'.")
+
+    return constraint_class.load_constraint_from_dict(parameters=parameters)

{sdv-1.36.3.dev0 → sdv-1.37.0}/sdv/cag/base.py

@@ -59,6 +59,37 @@ class BaseConstraint:
         args_string = ', '.join(custom_args)
         return f'{class_name}({args_string})'
 
+    def get_constraint_dict(self):
+        """Return the constraint as a serialiazable dict.
+
+        Returns:
+            dict:
+                A dictionary with the following keys:
+                    - `class_name` [str]: The name of the constraint class.
+                    - `parameters` [dict]: A dictionary of the init parameters used to
+                       create this constraint instance.
+        """
+        args = inspect.getfullargspec(self.__init__)
+        keys = args.args[1:]
+        instanced = {}
+        for key in keys:
+            if hasattr(self, key) or hasattr(self, f'_{key}'):
+                instanced[key] = getattr(self, key, getattr(self, f'_{key}', None))
+        missing_attrs = list(set(keys) - set(instanced.keys()))
+        if missing_attrs:
+            missing_attrs = sorted(missing_attrs)
+            raise AttributeError(
+                'Cannot convert constraint to dictionary because required parameters '
+                f'{missing_attrs} are not saved as attributes on the constraint.'
+            )
+
+        return {'class_name': self.__class__.__name__, 'parameters': instanced}
+
+    @classmethod
+    def load_constraint_from_dict(cls, parameters):
+        """Uses the given parameters to recreate an instance of the constraint."""
+        return cls(**parameters)
+
     def __init__(self):
         self.metadata = None
         self._fitted = False

{sdv-1.36.3.dev0 → sdv-1.37.0}/sdv/cag/programmable_constraint.py

@@ -1,5 +1,6 @@
 """Programmable constraints base classes."""
 
+import inspect
 from copy import deepcopy
 
 from sdv.cag.base import BaseConstraint
@@ -10,6 +11,11 @@ class ProgrammableConstraint:
 
     _is_single_table = True
 
+    @classmethod
+    def load_constraint_from_dict(cls, parameters):
+        """Uses the given parameters to recreate an instance of the constraint."""
+        return cls(**parameters)
+
     def validate(self, metadata):
         """Validates that the metadata is compatible with the constraint and its parameters.
 
@@ -133,6 +139,36 @@ class ProgrammableConstraintHarness(BaseConstraint):
         self.table_name = getattr(self.programmable_constraint, 'table_name', None)
         self._is_single_table = self.programmable_constraint._is_single_table
 
+    def get_constraint_dict(self):
+        """Return the constraint as a serialiazable dict.
+
+        Returns:
+            dict:
+                A dictionary with the following keys:
+                    - `class_name` [str]: The name of the constraint class.
+                    - `parameters` [dict]: A dictionary of the init parameters used to
+                       create this constraint instance.
+        """
+        args = inspect.getfullargspec(self.programmable_constraint.__init__)
+        keys = args.args[1:]
+        instanced = {}
+        constraint = self.programmable_constraint
+        for key in keys:
+            if hasattr(constraint, key) or hasattr(constraint, f'_{key}'):
+                instanced[key] = getattr(constraint, key, getattr(constraint, f'_{key}', None))
+        missing_attrs = list(set(keys) - set(instanced.keys()))
+        if missing_attrs:
+            missing_attrs = sorted(missing_attrs)
+            raise AttributeError(
+                'Cannot convert constraint to dictionary because required parameters '
+                f'{missing_attrs} are not saved as attributes on the constraint.'
+            )
+
+        return {
+            'class_name': self.programmable_constraint.__class__.__name__,
+            'parameters': instanced,
+        }
+
     def _validate_constraint_with_metadata(self, metadata):
         self.programmable_constraint.validate(metadata)
 

{sdv-1.36.3.dev0 → sdv-1.37.0}/sdv/datasets/demo.py

@@ -55,6 +55,20 @@ def _create_s3_client(bucket, credentials=None):
 
 
 def _get_data_from_bucket(object_key, bucket, client):
+    """Get a file from an S3 bucket as a bytes object.
+
+    Args:
+        object_key (str):
+            The key of the object to get.
+        bucket (str):
+            The name of the bucket to get the object from.
+        client (botocore.client.S3):
+            S3 client.
+
+    Returns:
+        bytes:
+            The file data from the S3 object as a bytes object.
+    """
     response = client.get_object(Bucket=bucket, Key=object_key)
     return response['Body'].read()
 
@@ -271,11 +285,31 @@ def handle_aws_client_errors(error_message_builder):
 
 
 def _download(modality, dataset_name, bucket, credentials=None):
-    """Download dataset resources from a bucket.
+    """Download dataset resources from S3 bucket and return the bytes.
+
+    Args:
+        modality (str):
+            The modality of the dataset: ``'single_table'``, ``'multi_table'``,
+            ``'sequential'``.
+        dataset_name (str):
+            The name of the dataset.
+        bucket (str):
+            The name of the bucket to download from.
+        credentials (dict or None):
+            Dictionary containing DataCebo license key and username. It takes the form:
+            { 'username': 'example@datacebo.com', 'license_key': '<MY_LICENSE_KEY>' }
+            If None, the function will use the default credentials.
 
     Returns:
-        tuple:
-            (BytesIO(zip_bytes), metadata_bytes)
+        tuple[BytesIO, bytes]:
+            (data_bytes, metadata_bytes)
+            The data is bytes of the ``data.zip`` and
+            ``metadata_bytes`` is the raw bytes of the metadata JSON.
+
+    Raises:
+        DemoResourceNotFoundError:
+            If the dataset prefix is missing in the bucket, if ``data.zip`` is
+            missing, or if no V1 metadata file is present.
     """
     client = _create_s3_client(bucket=bucket, credentials=credentials)
     dataset_prefix = f'{modality}/{dataset_name}/'
@@ -286,86 +320,64 @@ def _download(modality, dataset_name, bucket, credentials=None):
     )
     contents = _list_objects(dataset_prefix, bucket=bucket, client=client)
     zip_key = _find_data_zip_key(contents, dataset_prefix, bucket)
-    zip_bytes = _get_data_from_bucket(zip_key, bucket=bucket, client=client)
+    data_bytes = io.BytesIO(_get_data_from_bucket(zip_key, bucket=bucket, client=client))
     metadata_bytes = _get_first_v1_metadata_bytes(
         contents, dataset_prefix, bucket=bucket, client=client
     )
 
-    return io.BytesIO(zip_bytes), metadata_bytes
+    return data_bytes, metadata_bytes
 
 
-def _extract_data(bytes_io, output_folder_name):
-    with ZipFile(bytes_io) as zf:
-        if output_folder_name:
-            os.makedirs(output_folder_name, exist_ok=True)
-            zf.extractall(output_folder_name)
-
-        else:
-            in_memory_directory = {}
-            for name in zf.namelist():
-                in_memory_directory[name] = zf.read(name)
+def _load_data_from_zip(zip_bytes, bucket, dataset_name, output_folder_name=None):
+    """Load CSV tables from in-memory zip bytes into a dict of DataFrames.
 
-            return in_memory_directory
+    This function iterates over the zip bytes and parses each CSV entry with
+    ``pandas.read_csv``. Non-CSV entries are recorded as skipped. When
+    ``output_folder_name`` is provided, the archive is also extracted to disk
+    as a side effect so the caller keeps a local copy.
 
+    Args:
+        zip_bytes (io.BytesIO):
+            File-like object containing the bytes of ``data.zip``.
+        bucket (str):
+            The name of the bucket the zip was downloaded from. Used only for
+            error messages.
+        dataset_name (str):
+            The name of the dataset. Used only for error messages.
+        output_folder_name (str or None):
+            Optional folder path where the zip will also be extracted to disk so
+            the user keeps a local copy. The returned data dict is always built
+            in-memory. If ``None``, no folder is created.
 
-def _get_data_with_output_folder(output_folder_name):
-    """Load CSV tables from an extracted folder on disk.
+    Returns:
+        dict[str, pandas.DataFrame]:
+            Mapping of table name to DataFrame.
 
-    Returns a tuple of (data_dict, skipped_files).
-    Non-CSV files are ignored.
+    Raises:
+        DemoResourceNotFoundError:
+            If the zip contains no valid CSV entries.
     """
     data = {}
     skipped_files = []
-    for root, _dirs, files in os.walk(output_folder_name):
-        for filename in files:
+    with ZipFile(zip_bytes, 'r') as z:
+        if output_folder_name:
+            os.makedirs(output_folder_name, exist_ok=True)
+            z.extractall(output_folder_name)
+
+        for filename in z.namelist():
             if not filename.lower().endswith('.csv'):
                 skipped_files.append(filename)
                 continue
 
             table_name = Path(filename).stem
-            data_path = os.path.join(root, filename)
             try:
-                data[table_name] = pd.read_csv(data_path)
+                with z.open(filename) as f:
+                    data[table_name] = pd.read_csv(f, low_memory=False)
             except UnicodeDecodeError:
-                data[table_name] = pd.read_csv(data_path, encoding=FALLBACK_ENCODING)
+                with z.open(filename) as f:
+                    data[table_name] = pd.read_csv(f, low_memory=False, encoding=FALLBACK_ENCODING)
             except Exception as e:
-                rel = os.path.relpath(data_path, output_folder_name)
-                skipped_files.append(f'{rel}: {e}')
-
-    return data, skipped_files
-
-
-def _get_data_without_output_folder(in_memory_directory):
-    """Load CSV tables directly from in-memory zip contents.
-
-    Returns a tuple of (data_dict, skipped_files).
-    Non-CSV entries are ignored.
-    """
-    data = {}
-    skipped_files = []
-    for filename, file_ in in_memory_directory.items():
-        if not filename.lower().endswith('.csv'):
-            skipped_files.append(filename)
-            continue
-
-        table_name = Path(filename).stem
-        try:
-            data[table_name] = pd.read_csv(io.BytesIO(file_), low_memory=False)
-        except UnicodeDecodeError:
-            data[table_name] = pd.read_csv(
-                io.BytesIO(file_), low_memory=False, encoding=FALLBACK_ENCODING
-            )
-        except Exception as e:
-            skipped_files.append(f'{filename}: {e}')
-
-    return data, skipped_files
-
-
-def _get_data(modality, output_folder_name, in_memory_directory, bucket, dataset_name):
-    if output_folder_name:
-        data, skipped_files = _get_data_with_output_folder(output_folder_name)
-    else:
-        data, skipped_files = _get_data_without_output_folder(in_memory_directory)
+                skipped_files.append(f'{filename}: {e}')
 
     if skipped_files:
         warnings.warn('Skipped files: ' + ', '.join(sorted(skipped_files)))
@@ -376,9 +388,6 @@ def _get_data(modality, output_folder_name, in_memory_directory, bucket, dataset
             'The dataset is missing `csv` file/s.'
         )
 
-    if modality != 'multi_table':
-        data = data.popitem()[1]
-
     return data
 
 
@@ -464,17 +473,18 @@ def download_demo(
     _validate_modalities(modality)
     _validate_output_folder(output_folder_name)
 
-    data_io, metadata_bytes = _download(modality, dataset_name, s3_bucket_name, credentials)
-    in_memory_directory = _extract_data(data_io, output_folder_name)
-    data = _get_data(
+    data, metadata_bytes = _download(
         modality,
-        output_folder_name,
-        in_memory_directory,
-        s3_bucket_name,
         dataset_name,
+        s3_bucket_name,
+        credentials,
     )
-    metadata = _get_metadata(metadata_bytes, dataset_name, output_folder_name)
+    data = _load_data_from_zip(data, s3_bucket_name, dataset_name, output_folder_name)
 
+    if modality != 'multi_table':
+        data = data.popitem()[1]
+
+    metadata = _get_metadata(metadata_bytes, dataset_name, output_folder_name)
     return data, metadata
 
 

{sdv-1.36.3.dev0 → sdv-1.37.0}/sdv/multi_table/base.py

@@ -3,10 +3,13 @@
 import contextlib
 import datetime
 import inspect
+import json
 import operator
+import traceback
 import warnings
 from collections import defaultdict
 from copy import deepcopy
+from pathlib import Path
 
 import cloudpickle
 import numpy as np
@@ -25,6 +28,7 @@ from sdv.cag._utils import (
     _convert_to_snake_case,
     _get_invalid_rows,
     _validate_constraints,
+    load_constraint_from_dict,
 )
 from sdv.cag.programmable_constraint import ProgrammableConstraint, ProgrammableConstraintHarness
 from sdv.errors import (
@@ -269,8 +273,7 @@ class BaseMultiTableSynthesizer:
 
         self._single_table_constraints += single_table_constraints
 
-    def get_constraints(self):
-        """Get a copy of the list of constraints applied to the synthesizer."""
+    def _get_all_constraints_list(self):
         if not hasattr(self, 'constraints'):
             return []
         constraints = []
@@ -285,6 +288,68 @@ class BaseMultiTableSynthesizer:
 
         return constraints
 
+    def get_constraints(self, filepath=None):
+        """Get a list of constraint-augmented generation constraints applied to the synthesizer.
+
+        If `filepath` is provided, will save the constraints in JSON format to `filepath`.
+        Otherwise, will return a list of the constraints applied to the synthesizer.
+
+        Args:
+            filepath (str, optional):
+                Path where the constraints applied to the synthesizer will be serialized. If `None`,
+                will return a list of the applied constraints instead. Defaults to `None`.
+        """
+        constraints = self._get_all_constraints_list()
+        if filepath is None:
+            return constraints
+
+        path = Path(filepath)
+        if path.exists():
+            raise ValueError(
+                f"Cannot save constraints to file because '{filepath}' already exists."
+            )
+
+        constraints_dict_list = [constraint.get_constraint_dict() for constraint in constraints]
+        with open(path, 'w') as file:
+            json.dump(constraints_dict_list, file, indent=4)
+
+    def set_constraints(self, filepath):
+        """Add all the constraints in the file to the synthesizer.
+
+        If any constraints have been added to the synthesizer, they will be removed before
+        the constraints from the file are set.
+
+        Args:
+            filepath (str):
+                The string path to the file containing the constraints to set on the synthesizer.
+        """
+        if self.get_constraints():
+            raise SynthesizerInputError(
+                'Cannot `set_constraints` since constraints have already been applied.'
+            )
+
+        with open(filepath, 'r') as f:
+            constraints_json = json.load(f)
+
+        constraint_list = []
+        for constraint_dict in constraints_json:
+            try:
+                constraint_list.append(load_constraint_from_dict(constraint_dict))
+            except Exception as e:
+                warnings.warn(
+                    f'Could not load constraint ({constraint_dict}):\n'
+                    f'    {traceback.format_exception_only(type(e), e)[0]}'
+                )
+
+        for constraint in constraint_list:
+            try:
+                self.add_constraints([constraint])
+            except Exception as e:
+                warnings.warn(
+                    f'Could not add constraint ({constraint}):\n'
+                    f'    {traceback.format_exception_only(type(e), e)[0]}'
+                )
+
     def validate_constraints(self, synthetic_data):
         """Validate synthetic_data against the constraints.
 

{sdv-1.36.3.dev0 → sdv-1.37.0}/sdv/single_table/base.py

@@ -3,14 +3,17 @@
 import datetime
 import functools
 import inspect
+import json
 import logging
 import math
 import operator
 import os
+import traceback
 import uuid
 import warnings
 from collections import defaultdict
 from copy import deepcopy
+from pathlib import Path
 
 import cloudpickle
 import copulas
@@ -35,6 +38,7 @@ from sdv.cag._utils import (
     _convert_to_snake_case,
     _get_invalid_rows,
     _validate_constraints_single_table,
+    load_constraint_from_dict,
 )
 from sdv.cag.programmable_constraint import ProgrammableConstraint, ProgrammableConstraintHarness
 from sdv.data_processing.data_processor import DataProcessor
@@ -497,8 +501,7 @@ class BaseSynthesizer:
             locales=self.locales,
         )
 
-    def get_constraints(self):
-        """Get a list of constraint-augmented generation constraints applied to the synthesizer."""
+    def _get_all_constraints_list(self):
         constraints = []
         for constraint in self._chained_constraints + self._reject_sampling_constraints:
             if isinstance(constraint, ProgrammableConstraintHarness):
@@ -508,6 +511,69 @@ class BaseSynthesizer:
 
         return constraints
 
+    def get_constraints(self, filepath=None):
+        """Get a list of constraint-augmented generation constraints applied to the synthesizer.
+
+        If `filepath` is provided, will save the constraints in JSON format to `filepath`.
+        Otherwise, will return a list of the constraints applied to the synthesizer.
+
+        Args:
+            filepath (str, optional):
+                Path where the constraints applied to the synthesizer will be serialized. If `None`,
+                will return a list of the applied constraints instead. Defaults to `None`.
+        """
+        constraints = self._get_all_constraints_list()
+        if filepath is None:
+            return constraints
+
+        path = Path(filepath)
+        if path.exists():
+            raise ValueError(
+                f"Cannot save constraints to file because '{filepath}' already exists."
+            )
+
+        constraints_dict_list = [constraint.get_constraint_dict() for constraint in constraints]
+
+        with open(path, 'w') as file:
+            json.dump(constraints_dict_list, file, indent=4)
+
+    def set_constraints(self, filepath):
+        """Add all the constraints in the file to the synthesizer.
+
+        If any constraints have been added to the synthesizer, they will be removed before
+        the constraints from the file are set.
+
+        Args:
+            filepath (str):
+                The string path to the file containing the constraints to set on the synthesizer.
+        """
+        if self.get_constraints():
+            raise SynthesizerInputError(
+                'Cannot `set_constraints` since constraints have already been applied.'
+            )
+
+        with open(filepath, 'r') as f:
+            constraints_json = json.load(f)
+
+        constraint_list = []
+        for constraint_dict in constraints_json:
+            try:
+                constraint_list.append(load_constraint_from_dict(constraint_dict))
+            except Exception as e:
+                warnings.warn(
+                    f'Could not load constraint ({constraint_dict}):\n'
+                    f'    {traceback.format_exception_only(type(e), e)[0]}'
+                )
+
+        for constraint in constraint_list:
+            try:
+                self.add_constraints([constraint])
+            except Exception as e:
+                warnings.warn(
+                    f'Could not add constraint ({constraint}):\n'
+                    f'    {traceback.format_exception_only(type(e), e)[0]}'
+                )
+
     def validate_constraints(self, synthetic_data):
         """Validate synthetic_data against the constraints.
 

{sdv-1.36.3.dev0 → sdv-1.37.0/sdv.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sdv
-Version: 1.36.3.dev0
+Version: 1.37.0
 Summary: Generate synthetic data for single table, multi table and sequential data
 Author-email: "DataCebo, Inc." <info@sdv.dev>
 License-Expression: BUSL-1.1