squirrels 0.1.0__py3-none-any.whl → 0.6.0.post0__py3-none-any.whl

squirrels/_parameter_options.py

@@ -0,0 +1,348 @@
+from typing import TypeVar, Iterable, Any
+from dataclasses import dataclass
+from decimal import Decimal, InvalidOperation as InvalidDecimalConversion
+from datetime import datetime, date
+from abc import ABCMeta, abstractmethod
+
+from ._utils import ConfigurationError
+
+Number = Decimal | int | float | str
+Comparables = TypeVar("Comparables", Decimal, date)
+
+
+@dataclass
+class ParameterOption(metaclass=ABCMeta):
+    """
+    Abstract class for parameter options
+    """
+    _user_groups: frozenset[Any]
+    _parent_option_ids: frozenset[str]
+
+    @abstractmethod
+    def __init__(
+        self, *, user_groups: Iterable[Any] | str = frozenset(), parent_option_ids: Iterable[str] | str = frozenset(), **kwargs
+    ) -> None:
+        self._user_groups = frozenset({user_groups} if isinstance(user_groups, str) else user_groups)
+        self._parent_option_ids = frozenset({parent_option_ids} if isinstance(parent_option_ids, str) else parent_option_ids)
+
+    def _validate_lower_upper_values(self, lower_label: str, lower_value: Comparables, upper_label: str, upper_value: Comparables):
+        if lower_value > upper_value:
+            raise ConfigurationError(f'The {lower_label} "{lower_value}" must be less than or equal to the {upper_label} "{upper_value}"')
+
+    def _is_valid(self, user_group: Any, selected_parent_option_ids: Iterable[str] | None) -> bool:
+        """
+        Checks if this option is valid given the selected parent options and user group of user if applicable.
+        
+        Arguments:
+            user_group: The value of the user's "user group attribute". Only None when "user_attribute" is not specified
+                for the Parameter factory.
+            selected_parent_option_ids: List of selected option ids from the parent parameter. Only None when the Parameter
+                object has no parent parameter.
+        
+        Returns:
+            True if valid, False otherwise
+        """
+        if user_group is not None and user_group not in self._user_groups:
+            return False
+
+        if selected_parent_option_ids is not None and self._parent_option_ids.isdisjoint(selected_parent_option_ids):
+            return False
+        
+        return True
+    
+
+@dataclass
+class SelectParameterOption(ParameterOption):
+    """
+    Parameter option for a select parameter
+    """
+    _identifier: str
+    _label: str
+    _is_default: bool
+    custom_fields: dict[str, Any]
+
+    def __init__(
+        self, id: str, label: str, *, is_default: bool = False, user_groups: Iterable[Any] | str = frozenset(), 
+        parent_option_ids: Iterable[str] | str = frozenset(), custom_fields: dict[str, Any] = {}, **kwargs
+    ) -> None:
+        """
+        Constructor for SelectParameterOption
+
+        Arguments:
+            identifier: Unique identifier for this option that never changes over time
+            label: Human readable label that gets shown as a dropdown option
+            is_default: True if this is a default option, False otherwise
+            user_groups: The user groups this parameter option would show for if "user_attribute" is specified in the Parameter factory
+            parent_option_ids: Set of parent option ids this parameter option would show for if "parent" is specified in the Parameter factory
+            custom_fields: Dictionary to associate custom attributes to the parameter option
+            **kwargs: Any additional keyword arguments specified (except the ones above) gets included into custom_fields as well
+        """
+        super().__init__(user_groups=user_groups, parent_option_ids=parent_option_ids)
+        self._identifier = id
+        self._label = label
+        self._is_default = is_default
+        self.custom_fields = {
+            **kwargs, **custom_fields, **self._to_json_dict()
+        }
+
+    def get_custom_field(self, field: str, *, default_field: str | None = None, default: Any = None, **kwargs) -> Any:
+        """
+        Get field value from the custom_fields attribute
+
+        Arguments:
+            field: The key to use to fetch the custom field from "custom_fields"
+            default_field: If value at "field" key does not exist in "custom_fields", then this is used instead as the field (if not None)
+            default: If value at "field" or "default_field" (if not None) key does not exist in "custom_fields", then this value 
+                is used as default, or throws an error if None
+        
+        Returns:
+            The type of the custom field
+        """
+        if default_field is not None:
+            default = self.get_custom_field(default_field, default=default)
+        
+        if default is not None:
+            selected_field = self.custom_fields.get(field, default)
+        else:
+            try:
+                selected_field = self.custom_fields[field]
+            except KeyError as e:
+                raise ConfigurationError(f"Field '{field}' must exist for parameter option {self._to_json_dict()}") from e
+        
+        return selected_field
+    
+    def _to_json_dict(self):
+        return {'id': self._identifier, 'label': self._label}
+
+
+@dataclass
+class _DateTypeParameterOption(ParameterOption):
+    """
+    Abstract class (or type) for date type parameter options
+    """
+    _min_date: date | None
+    _max_date: date | None
+    _date_format: str
+
+    @abstractmethod
+    def __init__(
+        self, *, min_date: str | date | None = None, max_date: str | date | None = None, date_format: str = '%Y-%m-%d', 
+        user_groups: Iterable[Any] | str = frozenset(), parent_option_ids: Iterable[str] | str = frozenset(), **kwargs
+    ) -> None:
+        super().__init__(user_groups=user_groups, parent_option_ids=parent_option_ids)
+        self._date_format = date_format
+        self._min_date, self._max_date = None, None # preset for using _validate_date()
+        self._min_date = self._validate_date(min_date) if min_date is not None else None
+        self._max_date = self._validate_date(max_date) if max_date is not None else None
+        if self._min_date is not None and self._max_date is not None:
+            self._validate_lower_upper_values("min_date", self._min_date, "max_date", self._max_date)
+
+    def _validate_date(self, date_str: str | date) -> date:
+        try:
+            date_obj = datetime.strptime(date_str, self._date_format).date() if isinstance(date_str, str) else date_str
+        except ValueError as e:
+            raise ConfigurationError(f'Invalid format for date "{date_str}".') from e
+        
+        if self._min_date is not None and date_obj < self._min_date:
+            raise ConfigurationError(f'The provided date "{date_obj}" is less than the min date "{self._min_date}"')
+        if self._max_date is not None and date_obj > self._max_date:
+            raise ConfigurationError(f'The provided date "{date_obj}" is greater than the max date "{self._max_date}"')
+        
+        return date_obj
+    
+
+@dataclass
+class DateParameterOption(_DateTypeParameterOption):
+    """
+    Parameter option for default dates if it varies based on selection of another parameter
+    """
+    _default_date: date
+
+    def __init__(
+        self, default_date: str | date, *, min_date: str | date | None = None, max_date: str | date | None = None, date_format: str = '%Y-%m-%d',
+        user_groups: Iterable[Any] | str = frozenset(), parent_option_ids: Iterable[str] | str = frozenset(), **kwargs
+    ) -> None:
+        """
+        Constructor for DateParameterOption
+
+        Arguments:
+            default_date: Default date for this option
+            min_date: Minimum date for this option
+            max_date: Maximum date for this option
+            date_format: Format of the default date, default is '%Y-%m-%d'
+            user_groups: The user groups this parameter option would show for if "user_attribute" is specified in the Parameter factory
+            parent_option_ids: Set of parent option ids this parameter option would show for if "parent" is specified in the Parameter factory
+        """
+        super().__init__(
+            date_format=date_format, min_date=min_date, max_date=max_date, user_groups=user_groups, parent_option_ids=parent_option_ids
+        )
+        self._default_date = self._validate_date(default_date)
+
+
+@dataclass
+class DateRangeParameterOption(_DateTypeParameterOption):
+    """
+    Parameter option for default dates if it varies based on selection of another parameter
+    """
+    _default_start_date: date
+    _default_end_date: date
+
+    def __init__(
+        self, default_start_date: str | date, default_end_date: str | date, *, min_date: str | date | None = None, 
+        max_date: str | date | None = None, date_format: str = '%Y-%m-%d', user_groups: Iterable[Any] | str = frozenset(), 
+        parent_option_ids: Iterable[str] | str = frozenset(), **kwargs
+    ) -> None:
+        """
+        Constructor for DateRangeParameterOption
+
+        Arguments:
+            default_start_date: Default start date for this option
+            default_end_date: Default end date for this option
+            min_date: Minimum date for this option
+            max_date: Maximum date for this option
+            date_format: Format of the default date, default is '%Y-%m-%d'
+            user_groups: The user groups this parameter option would show for if "user_attribute" is specified in the Parameter factory
+            parent_option_ids: Set of parent option ids this parameter option would show for if "parent" is specified in the Parameter factory
+        """
+        super().__init__(
+            date_format=date_format, min_date=min_date, max_date=max_date, user_groups=user_groups, parent_option_ids=parent_option_ids
+        )
+        self._default_start_date = self._validate_date(default_start_date)
+        self._default_end_date = self._validate_date(default_end_date)
+        self._validate_lower_upper_values("default_start_date", self._default_start_date, "default_end_date", self._default_end_date)
+
+
+@dataclass
+class _NumericParameterOption(ParameterOption):
+    """
+    Abstract class (or type) for numeric parameter options
+    """
+    _min_value: Decimal
+    _max_value: Decimal
+    _increment: Decimal
+
+    @abstractmethod
+    def __init__(
+        self, min_value: Number, max_value: Number, *, increment: Number = 1, user_groups: Iterable[Any] | str = frozenset(), 
+        parent_option_ids: Iterable[str] | str = frozenset(), **kwargs
+    ) -> None:
+        super().__init__(user_groups=user_groups, parent_option_ids=parent_option_ids)
+        try:
+            self._min_value = Decimal(str(min_value))
+            self._max_value = Decimal(str(max_value))
+            self._increment = Decimal(str(increment))
+        except InvalidDecimalConversion as e:
+            raise ConfigurationError(f'Could not convert either min, max, or increment to number') from e
+        
+        self._validate_lower_upper_values("min_value", self._min_value, "max_value", self._max_value)
+
+        if (self._max_value - self._min_value) % self._increment != 0:
+            raise ConfigurationError(f'The increment "{self._increment}" must fit evenly between ' + 
+                f'the min_value "{self._min_value}" and max_value "{self._max_value}"')
+
+    def __value_in_range(self, value: Decimal) -> bool:
+        return self._min_value <= value <= self._max_value
+    
+    def __value_on_increment(self, value: Decimal) -> bool:
+        diff = (value - self._min_value)
+        return diff >= 0 and diff % self._increment == 0
+
+    def _validate_value(self, value: Number) -> Decimal:
+        try:
+            value = Decimal(str(value))
+        except InvalidDecimalConversion as e:
+            raise ConfigurationError(f'Could not convert "{value}" to number', e)
+        
+        if not self.__value_in_range(value):
+            raise ConfigurationError(f'The selected value "{value}" is outside of bounds ' +
+                f'"{self._min_value}" and "{self._max_value}".')
+        if not self.__value_on_increment(value):
+            raise ConfigurationError(f'The difference between selected value "{value}" and lower value ' +
+                f'"{self._min_value}" must be a multiple of increment "{self._increment}".')
+        return value
+
+
+@dataclass
+class NumberParameterOption(_NumericParameterOption):
+    """
+    Parameter option for default numbers if it varies based on selection of another parameter
+    """
+    _default_value: Decimal
+
+    def __init__(
+        self, min_value: Number, max_value: Number, *, increment: Number = 1, default_value: Number | None = None,
+        user_groups: Iterable[Any] | str = frozenset(), parent_option_ids: Iterable[str] | str = frozenset(), **kwargs
+    ) -> None:
+        """
+        Constructor for NumberParameterOption
+        
+        * Note that the "Number" type denotes an int, a Decimal (from decimal module), or a string that can be parsed to Decimal
+
+        Arguments:
+            min_value: Minimum selectable value
+            max_value: Maximum selectable value
+            increment: Increment of selectable values, and must fit evenly between min_value and max_value
+            default_value: Default value for this option, and must be selectable based on min_value, max_value, and increment
+            user_groups: The user groups this parameter option would show for if "user_attribute" is specified in the Parameter factory
+            parent_option_ids: Set of parent option ids this parameter option would show for if "parent" is specified in the Parameter factory
+        """
+        super().__init__(min_value, max_value, increment=increment, user_groups=user_groups, parent_option_ids=parent_option_ids)
+        self._default_value = self._validate_value(default_value) if default_value is not None else self._min_value
+
+
+@dataclass
+class NumberRangeParameterOption(_NumericParameterOption):
+    """
+    Parameter option for default numeric ranges if it varies based on selection of another parameter
+    """
+    _default_lower_value: Decimal
+    _default_upper_value: Decimal
+
+    def __init__(
+        self, min_value: Number, max_value: Number, *, increment: Number = 1, default_lower_value: Number | None = None, 
+        default_upper_value: Number | None = None, user_groups: Iterable[Any] | str = frozenset(), 
+        parent_option_ids: Iterable[str] | str = frozenset(), **kwargs
+    ) -> None:
+        """
+        Constructor for NumberRangeParameterOption
+        
+        * Note that the "Number" type denotes an int, a Decimal (from decimal module), or a string that can be parsed to Decimal
+
+        Arguments:
+            min_value: Minimum selectable value
+            max_value: Maximum selectable value
+            increment: Increment of selectable values, and must fit evenly between min_value and max_value
+            default_lower_value: Default lower value for this option, and must be selectable based on min_value, max_value, and increment
+            default_upper_value: Default upper value for this option, and must be selectable based on min_value, max_value, and increment. 
+                    Must also be greater than default_lower_value
+            user_groups: The user groups this parameter option would show for if "user_attribute" is specified in the Parameter factory
+            parent_option_ids: Set of parent option ids this parameter option would show for if "parent" is specified in the Parameter factory
+        """
+        super().__init__(min_value, max_value, increment=increment, user_groups=user_groups, parent_option_ids=parent_option_ids)
+        self._default_lower_value = self._validate_value(default_lower_value) if default_lower_value is not None else self._min_value
+        self._default_upper_value = self._validate_value(default_upper_value) if default_upper_value is not None else self._max_value
+        self._validate_lower_upper_values("default_lower_value", self._default_lower_value, "default_upper_value", self._default_upper_value)
+
+
+@dataclass
+class TextParameterOption(ParameterOption):
+    """
+    Parameter option for default text values if it varies based on selection of another parameter
+    """
+    _default_text: str
+
+    def __init__(
+        self, *, default_text: str = "", user_groups: Iterable[Any] | str = frozenset(), 
+        parent_option_ids: Iterable[str] | str = frozenset(), **kwargs
+    ) -> None:
+        """
+        Constructor for TextParameterOption
+
+        Arguments:
+            default_text: Default text for this option
+            user_groups: The user groups this parameter option would show for if "user_attribute" is specified in the Parameter factory
+            parent_option_ids: Set of parent option ids this parameter option would show for if "parent" is specified in the Parameter factory
+        """
+        super().__init__(user_groups=user_groups, parent_option_ids=parent_option_ids)
+        self._default_text = default_text
+    

squirrels/_parameter_sets.py

@@ -0,0 +1,207 @@
+from __future__ import annotations
+from typing import Optional, Sequence, Callable, Any
+from dataclasses import dataclass, field
+from collections import OrderedDict
+import time, concurrent.futures, polars as pl
+
+from . import _parameters as p, _utils as u, _constants as c, _parameter_configs as pc, _py_module as pm
+from ._schemas import response_models as rm
+from ._arguments.init_time_args import ParametersArgs
+from ._manifest import ParametersConfig, ManifestConfig
+from ._connection_set import ConnectionSet
+from ._seeds import Seeds
+from ._schemas.auth_models import AbstractUser
+from ._env_vars import SquirrelsEnvVars
+
+
+@dataclass
+class ParameterSet:
+    """
+    A wrapper class for a sequence of parameters with the selections applied as well
+    """
+    _parameters_dict: OrderedDict[str, p.Parameter]
+
+    def get_parameters_as_dict(self) -> dict[str, p.Parameter]:
+        return self._parameters_dict.copy()
+
+    def to_api_response_model0(self) -> rm.ParametersModel:
+        parameters = []
+        for x in self._parameters_dict.values():
+            parameters.append(x._to_api_response_model0())
+        return rm.ParametersModel(parameters=parameters)
+
+
+@dataclass
+class ParameterConfigsSet:
+    """
+    Pool of parameter configs, can create multiple for unit testing purposes
+    """
+    _data: dict[str, pc.ParameterConfigBase] = field(default_factory=OrderedDict)
+    _data_source_params: dict[str, pc.DataSourceParameterConfig] = field(default_factory=dict)
+        
+    def get(self, name: Optional[str]) -> Optional[pc.ParameterConfigBase]:
+        try:
+            return self._data[name] if name is not None else None
+        except KeyError as e:
+            raise u.ConfigurationError(f'Unable to find parameter named "{name}"') from e
+
+    def add(self, param_config: pc.ParameterConfigBase) -> None:
+        self._data[param_config.name] = param_config
+        if isinstance(param_config, pc.DataSourceParameterConfig):
+            self._data_source_params[param_config.name] = param_config
+    
+    def _get_all_ds_param_configs(self) -> Sequence[pc.DataSourceParameterConfig]:
+        return list(self._data_source_params.values())
+
+    def __convert_datasource_params(self, df_dict: dict[str, pl.DataFrame]) -> None:
+        done = set()
+        for curr_name in self._data_source_params:
+            stack = [curr_name] # Note: parents must be converted first before children
+            while stack:
+                name = stack[-1]
+                if name not in done:
+                    param = self._data_source_params.get(name, self.get(name))
+                    assert param is not None
+                    parent_name = param.parent_name
+                    if parent_name is not None and parent_name not in done:
+                        stack.append(parent_name)
+                        continue
+                    if isinstance(param, pc.DataSourceParameterConfig):
+                        if name not in df_dict:
+                            raise u.ConfigurationError(f'No reference data found for parameter "{name}"')
+                        self._data[name] = param.convert(df_dict[name])
+                    done.add(name)
+                stack.pop()
+    
+    def __validate_param_relationships(self) -> None:
+        for param_config in self._data.values():
+            assert isinstance(param_config, pc.ParameterConfig)
+            parent_name = param_config.parent_name
+            parent = self.get(parent_name)
+            if parent:
+                if not isinstance(param_config, pc.SelectionParameterConfig):
+                    if not isinstance(parent, pc.SingleSelectParameterConfig):
+                        raise u.ConfigurationError(f'Only single-select parameters can be parents of non-select parameters. ' +
+                                                   f'Parameter "{parent_name}" is the parent of non-select parameter ' +
+                                                   f'"{param_config.name}" but "{parent_name}" is not a single-select parameter.')
+                    seen = set()
+                    for option in param_config.all_options:
+                        lookup_keys = option._parent_option_ids
+                        if len(option._user_groups) > 0:
+                            lookup_keys = set((x, y) for x in option._parent_option_ids for y in option._user_groups)
+                        if not seen.isdisjoint(lookup_keys):
+                            raise u.ConfigurationError(f'Each distinct value of "parent option id" can only appear once (per user group)' +
+                                                       f'among the options of non-select parameter "{param_config.name}".')
+                        seen.update(lookup_keys)
+                
+                if not isinstance(parent, pc.SelectionParameterConfig):
+                    raise u.ConfigurationError(f'Only selection parameters can be parents. Parameter "{parent_name}" is the parent of ' +
+                                               f'"{param_config.name}" but "{parent_name}" is not a selection parameter.')
+                
+                parent._add_child_mutate(param_config)
+    
+    def _post_process_params(self, df_dict: dict[str, pl.DataFrame]) -> None:
+        self.__convert_datasource_params(df_dict)
+        self.__validate_param_relationships()
+    
+    def apply_selections(
+        self, dataset_params: Optional[Sequence[str]], selections: dict[str, Any], user: AbstractUser, *, parent_param: str | None = None
+    ) -> ParameterSet:
+        if dataset_params is None:
+            dataset_params = list(self._data.keys())
+        
+        parameters_by_name: dict[str, p.Parameter] = {}
+        params_to_process = [parent_param] if parent_param else dataset_params
+        params_to_process_set = set(params_to_process)
+        for some_name in params_to_process:
+            stack = [some_name] # Note: process parent selections first (if applicable) before children
+            while stack:
+                curr_name = stack[-1]
+                children = []
+                if curr_name not in parameters_by_name:
+                    param_conf = self.get(curr_name)
+                    assert isinstance(param_conf, pc.ParameterConfig)
+                    parent_name = param_conf.parent_name
+                    if parent_name is None:
+                        parent = None
+                    elif parent_name in params_to_process_set and parent_name not in parameters_by_name:
+                        stack.append(parent_name)
+                        continue
+                    else:
+                        parent = parameters_by_name.get(parent_name)
+                    assert isinstance(parent, p._SelectionParameter) or parent is None
+                    param = param_conf.with_selection(selections.get(curr_name), user, parent)
+                    parameters_by_name[curr_name] = param
+                    if isinstance(param_conf, pc.SelectionParameterConfig):
+                        children = list(x for x in param_conf.children.keys() if x in dataset_params)
+                stack.pop()
+                stack.extend(children)
+        
+        ordered_parameters = OrderedDict((key, parameters_by_name[key]) for key in dataset_params if key in parameters_by_name)
+        return ParameterSet(ordered_parameters)
+    
+    def get_all_api_field_info(self) -> dict[str, pc.APIParamFieldInfo]:
+        api_field_infos = {}
+        for param, config in self._data.items():
+            assert isinstance(config, pc.ParameterConfig)
+            api_field_infos[param] = config.get_api_field_info()
+        return api_field_infos
+
+
+class ParameterConfigsSetIO:
+    """
+    Static class for the singleton object of ParameterConfigsSet
+    """
+    param_factories: list[Callable[[ParametersArgs], pc.ParameterConfigBase]] = []  # this is static (set in load_from_file) to stage the functions from pyconfigs/parameters.py before using them
+    
+    @classmethod
+    def _get_df_dict_from_data_sources(
+        cls, param_configs_set: ParameterConfigsSet, default_conn_name: str, seeds: Seeds, conn_set: ConnectionSet, datalake_db_path: str
+    ) -> dict[str, pl.DataFrame]:
+        
+        def get_dataframe(ds_param_config: pc.DataSourceParameterConfig) -> tuple[str, pl.DataFrame]:
+            return ds_param_config.name, ds_param_config.get_dataframe(default_conn_name, conn_set, seeds, datalake_db_path)
+        
+        ds_param_configs = param_configs_set._get_all_ds_param_configs()
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            df_dict = dict(executor.map(get_dataframe, ds_param_configs))
+        
+        return df_dict
+    
+    @classmethod
+    def _add_from_dict(cls, param_configs_set: ParameterConfigsSet, param_config: ParametersConfig) -> None:
+        ptype = getattr(p, param_config.type)
+        factory = getattr(ptype, param_config.factory)
+        obj = factory(**param_config.arguments)
+        param_configs_set.add(obj)
+    
+    @classmethod
+    def load_from_file(
+        cls, logger: u.Logger, env_vars: SquirrelsEnvVars, manifest_cfg: ManifestConfig, seeds: Seeds, conn_set: ConnectionSet, 
+        param_args: ParametersArgs
+    ) -> ParameterConfigsSet:
+        start = time.time()
+        param_configs_set = ParameterConfigsSet()
+
+        for param_as_dict in manifest_cfg.parameters:
+            cls._add_from_dict(param_configs_set, param_as_dict)
+        
+        # adds to cls.param_factories as side effect
+        main_result = pm.run_pyconfig_main(env_vars.project_path, c.PARAMETERS_FILE, {"sqrl": param_args})  
+        param_factories = cls.param_factories
+        cls.param_factories = []
+
+        for param_factory in param_factories:
+            param_configs_set.add(param_factory(param_args))
+        
+        if isinstance(main_result, list):
+            for param_config in main_result:
+                param_configs_set.add(param_config)
+        
+        default_conn_name = env_vars.connections_default_name_used
+        datalake_db_path = env_vars.vdl_catalog_db_path
+        df_dict = cls._get_df_dict_from_data_sources(param_configs_set, default_conn_name, seeds, conn_set, datalake_db_path)
+        param_configs_set._post_process_params(df_dict)
+        
+        logger.log_activity_time("loading parameters", start)
+        return param_configs_set