squirrels 0.5.0b2__py3-none-any.whl → 0.5.0b4__py3-none-any.whl

squirrels/_parameter_sets.py

@@ -1,13 +1,14 @@
 from __future__ import annotations
-from typing import Optional, Sequence, Any
+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 _utils as u, _constants as c, parameters as p, _parameter_configs as _pc, _py_module as pm, _api_response_models as arm
-from .arguments.init_time_args import ParametersArgs
+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, ConnectionsArgs
+from ._connection_set import ConnectionSet
 from ._seeds import Seeds
 from ._auth import BaseUser
 
@@ -22,11 +23,11 @@ class ParameterSet:
     def get_parameters_as_dict(self) -> dict[str, p.Parameter]:
         return self._parameters_dict.copy()
 
-    def to_api_response_model0(self) -> arm.ParametersModel:
+    def to_api_response_model0(self) -> rm.ParametersModel:
         parameters = []
         for x in self._parameters_dict.values():
             parameters.append(x._to_api_response_model0())
-        return arm.ParametersModel(parameters=parameters)
+        return rm.ParametersModel(parameters=parameters)
 
 
 @dataclass
@@ -34,21 +35,21 @@ 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)
+    _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]:
+    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:
+    def add(self, param_config: pc.ParameterConfigBase) -> None:
         self._data[param_config.name] = param_config
-        if isinstance(param_config, _pc.DataSourceParameterConfig):
+        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]:
+    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:
@@ -64,7 +65,7 @@ class ParameterConfigsSet:
                     if parent_name is not None and parent_name not in done:
                         stack.append(parent_name)
                         continue
-                    if isinstance(param, _pc.DataSourceParameterConfig):
+                    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])
@@ -73,12 +74,12 @@ class ParameterConfigsSet:
     
     def __validate_param_relationships(self) -> None:
         for param_config in self._data.values():
-            assert isinstance(param_config, _pc.ParameterConfig)
+            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):
+                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.')
@@ -92,7 +93,7 @@ class ParameterConfigsSet:
                                                        f'among the options of non-select parameter "{param_config.name}".')
                         seen.update(lookup_keys)
                 
-                if not isinstance(parent, _pc.SelectionParameterConfig):
+                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.')
                 
@@ -121,7 +122,7 @@ class ParameterConfigsSet:
                 children = []
                 if curr_name not in parameters_by_name:
                     param_conf = self.get(curr_name)
-                    assert isinstance(param_conf, _pc.ParameterConfig)
+                    assert isinstance(param_conf, pc.ParameterConfig)
                     parent_name = param_conf.parent_name
                     if parent_name is None:
                         parent = None
@@ -133,7 +134,7 @@ class ParameterConfigsSet:
                     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):
+                    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)
@@ -141,10 +142,10 @@ class ParameterConfigsSet:
         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]:
+    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)
+            assert isinstance(config, pc.ParameterConfig)
             api_field_infos[param] = config.get_api_field_info()
         return api_field_infos
 
@@ -153,44 +154,53 @@ class ParameterConfigsSetIO:
     """
     Static class for the singleton object of ParameterConfigsSet
     """
-    obj: ParameterConfigsSet # this is static (set in load_from_file) to simplify development experience for pyconfigs/parameters.py
+    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, default_conn_name: str, seeds: Seeds, conn_set: ConnectionSet) -> dict[str, pl.DataFrame]:
-        def get_dataframe(ds_param_config: _pc.DataSourceParameterConfig) -> tuple[str, pl.DataFrame]:
+    def _get_df_dict_from_data_sources(
+        cls, param_configs_set: ParameterConfigsSet, default_conn_name: str, seeds: Seeds, conn_set: ConnectionSet
+    ) -> 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)
         
-        ds_param_configs = cls.obj._get_all_ds_param_configs()
+        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_config: ParametersConfig) -> None:
+    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)
-        factory(**param_config.arguments)
-    
-    @classmethod
-    def get_param_args(cls, conn_args: ConnectionsArgs) -> ParametersArgs:
-        return ParametersArgs(conn_args.project_path, conn_args.proj_vars, conn_args.env_vars)
+        obj = factory(**param_config.arguments)
+        param_configs_set.add(obj)
     
     @classmethod
     def load_from_file(
         cls, logger: u.Logger, base_path: str, manifest_cfg: ManifestConfig, seeds: Seeds, conn_set: ConnectionSet, param_args: ParametersArgs
     ) -> ParameterConfigsSet:
         start = time.time()
-        cls.obj = ParameterConfigsSet()
+        param_configs_set = ParameterConfigsSet()
 
         for param_as_dict in manifest_cfg.parameters:
-            cls._add_from_dict(param_as_dict)
+            cls._add_from_dict(param_configs_set, param_as_dict)
+        
+        main_result = pm.run_pyconfig_main(base_path, c.PARAMETERS_FILE, {"sqrl": param_args})  # adds to cls.param_factories as side effect
+        param_factories = cls.param_factories
+        cls.param_factories = []
+
+        for param_factory in param_factories:
+            param_configs_set.add(param_factory(param_args))
         
-        pm.run_pyconfig_main(base_path, c.PARAMETERS_FILE, {"sqrl": param_args})
+        if isinstance(main_result, list):
+            for param_config in main_result:
+                param_configs_set.add(param_config)
         
         default_conn_name = manifest_cfg.env_vars.get(c.SQRL_CONNECTIONS_DEFAULT_NAME_USED, "default")
-        df_dict = cls._get_df_dict_from_data_sources(default_conn_name, seeds, conn_set)
-        cls.obj._post_process_params(df_dict)
+        df_dict = cls._get_df_dict_from_data_sources(param_configs_set, default_conn_name, seeds, conn_set)
+        param_configs_set._post_process_params(df_dict)
         
         logger.log_activity_time("loading parameters", start)
-        return cls.obj
+        return param_configs_set