squirrels 0.5.0b3__py3-none-any.whl → 0.6.0.post0__py3-none-any.whl

squirrels/_arguments/{_init_time_args.py → init_time_args.py}

@@ -4,60 +4,40 @@ import polars as pl
 
 from .. import _utils as u
 
+
 @dataclass
 class ConnectionsArgs:
     project_path: str
-    _proj_vars: dict[str, Any]
-    _env_vars: dict[str, str]
+    proj_vars: dict[str, Any]
+    env_vars: dict[str, str]
 
-    @property
-    def proj_vars(self) -> dict[str, Any]:
-        return self._proj_vars.copy()
-    
-    @property
-    def env_vars(self) -> dict[str, str]:
-        return self._env_vars.copy()
+    def __post_init__(self) -> None:
+        self.proj_vars = self.proj_vars.copy()
+        self.env_vars = self.env_vars.copy()
 
 
 @dataclass
-class ParametersArgs(ConnectionsArgs):
+class AuthProviderArgs(ConnectionsArgs):
     pass
 
 
 @dataclass
-class _WithConnectionDictArgs(ConnectionsArgs):
-    _connections: dict[str, Any]
-
-    @property
-    def connections(self) -> dict[str, Any]:
-        """
-        A dictionary of connection keys to SQLAlchemy Engines for database connections. 
-        
-        Can also be used to store other in-memory objects in advance such as ML models.
-        """
-        return self._connections.copy()
+class ParametersArgs(ConnectionsArgs):
+    pass
 
 
-class BuildModelArgs(_WithConnectionDictArgs):
+@dataclass
+class BuildModelArgs(ConnectionsArgs):
+    connections: dict[str, Any]
+    dependencies: Iterable[str]
+    _ref_func: Callable[[str], pl.LazyFrame]
+    _run_external_sql_func: Callable[[str, str], pl.DataFrame]
 
-    def __init__(
-        self, conn_args: ConnectionsArgs, _connections: dict[str, Any], 
-        dependencies: Iterable[str], 
-        ref: Callable[[str], pl.LazyFrame], 
-        run_external_sql: Callable[[str, str], pl.DataFrame]
-    ):
-        super().__init__(conn_args.project_path, conn_args.proj_vars, conn_args.env_vars, _connections)
-        self._dependencies = dependencies
-        self._ref = ref
-        self._run_external_sql = run_external_sql
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        self.connections = self.connections.copy()
+        self.dependencies = set(self.dependencies)
 
-    @property
-    def dependencies(self) -> set[str]:
-        """
-        The set of dependent data model names
-        """
-        return set(self._dependencies)
-    
     def ref(self, model: str) -> pl.LazyFrame:
         """
         Returns the result (as polars DataFrame) of a dependent model (predefined in "dependencies" function)
@@ -71,7 +51,7 @@ class BuildModelArgs(_WithConnectionDictArgs):
         Returns:
             A polars DataFrame
         """
-        return self._ref(model)
+        return self._ref_func(model)
 
     def run_external_sql(self, connection_name: str, sql_query: str, **kwargs) -> pl.DataFrame:
         """
@@ -84,20 +64,20 @@ class BuildModelArgs(_WithConnectionDictArgs):
         Returns:
             The query result as a polars DataFrame
         """
-        return self._run_external_sql(sql_query, connection_name)
+        return self._run_external_sql_func(sql_query, connection_name)
 
     def run_sql_on_dataframes(self, sql_query: str, *, dataframes: dict[str, pl.LazyFrame] | None = None, **kwargs) -> pl.DataFrame:
         """
         Uses a dictionary of dataframes to execute a SQL query in an embedded in-memory DuckDB database
 
         Arguments:
-            sql_query: The SQL query to run
+            sql_query: The SQL query to run (DuckDB dialect)
             dataframes: A dictionary of table names to their polars LazyFrame. If None, uses results of dependent models
         
         Returns:
             The result as a polars DataFrame from running the query
         """
         if dataframes is None:
-            dataframes = {x: self.ref(x) for x in self._dependencies}
+            dataframes = {x: self.ref(x) for x in self.dependencies}
 
         return u.run_sql_on_dataframes(sql_query, dataframes)

squirrels/_arguments/{_run_time_args.py → run_time_args.py}

@@ -1,45 +1,27 @@
 from typing import Callable, Any, Coroutine
+from dataclasses import dataclass, field, KW_ONLY
 import polars as pl
 
-from ._init_time_args import ParametersArgs, BuildModelArgs
-from .._auth import BaseUser
+from .init_time_args import ConnectionsArgs, ParametersArgs, BuildModelArgs
+from .._schemas.auth_models import AbstractUser
 from .._parameters import Parameter, TextValue
 
 
+@dataclass
 class ContextArgs(ParametersArgs):
-
-    def __init__(
-        self, param_args: ParametersArgs, 
-        user: BaseUser | None, 
-        prms: dict[str, Parameter], 
-        traits: dict[str, Any]
-    ):
-        super().__init__(param_args.project_path, param_args.proj_vars, param_args.env_vars)
-        self.user = user
-        self._prms = prms
-        self._traits = traits
-        self._placeholders = {}
-
-    @property
-    def prms(self) -> dict[str, Parameter]:
-        """
-        A dictionary of parameter names to parameter
-        """
-        return self._prms.copy()
-    
-    @property
-    def traits(self) -> dict[str, Any]:
-        """
-        A dictionary of dataset trait name to value
-        """
-        return self._traits.copy()
-
-    @property
-    def _placeholders_copy(self) -> dict[str, Any]:
-        """
-        A dictionary of placeholder name to placeholder value
-        """
-        return self._placeholders.copy()
+    user: AbstractUser
+    prms: dict[str, Parameter]
+    configurables: dict[str, str]
+    _conn_args: ConnectionsArgs
+    _: KW_ONLY
+    _placeholders: dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        self.prms = self.prms.copy()
+        self.configurables = self.configurables.copy()
+        self._conn_args = ConnectionsArgs(**self._conn_args.__dict__)
+        self._placeholders = self._placeholders.copy()
 
     def set_placeholder(self, placeholder: str, value: TextValue | Any) -> str:
         """
@@ -68,32 +50,14 @@ class ContextArgs(ParametersArgs):
         return (param_name in self.prms and self.prms[param_name].is_enabled())
 
 
-class ModelArgs(BuildModelArgs, ContextArgs):
-
-    def __init__(
-        self, ctx_args: ContextArgs, build_model_args: BuildModelArgs, 
-        ctx: dict[str, Any]
-    ):
-        super(ContextArgs, self).__init__(ctx_args.project_path, ctx_args.proj_vars, ctx_args.env_vars)
-        self._project_path = ctx_args.project_path
-        self._proj_vars = ctx_args.proj_vars
-        self._env_vars = ctx_args.env_vars
-        self.user = ctx_args.user
-        self._prms = ctx_args.prms
-        self._traits = ctx_args.traits
-        self._placeholders = ctx_args._placeholders_copy
-        self._connections = build_model_args.connections
-        self._dependencies = build_model_args.dependencies
-        self._ref = build_model_args.ref
-        self._run_external_sql = build_model_args.run_external_sql
-        self._ctx = ctx
-
-    @property
-    def ctx(self) -> dict[str, Any]:
-        """
-        Dictionary of context variables
-        """
-        return self._ctx.copy()
+@dataclass
+class ModelArgs(ContextArgs, BuildModelArgs):
+    ctx: dict[str, Any]
+
+    def __post_init__(self) -> None:
+        ContextArgs.__post_init__(self)
+        BuildModelArgs.__post_init__(self)
+        self.ctx = self.ctx.copy()
     
     def is_placeholder(self, placeholder: str) -> bool:
         """
@@ -122,14 +86,14 @@ class ModelArgs(BuildModelArgs, ContextArgs):
         return self._placeholders.get(placeholder)
 
 
-class DashboardArgs(ParametersArgs):
+@dataclass
+class DashboardArgs(ContextArgs):
+    ctx: dict[str, Any]
+    _get_dataset: Callable[[str, dict[str, Any]], Coroutine[Any, Any, pl.DataFrame]]
 
-    def __init__(
-        self, param_args: ParametersArgs, 
-        get_dataset: Callable[[str, dict[str, Any]], Coroutine[Any, Any, pl.DataFrame]]
-    ):
-        super().__init__(param_args.project_path, param_args.proj_vars, param_args.env_vars)
-        self._get_dataset = get_dataset
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        self.ctx = self.ctx.copy()
 
     async def dataset(self, name: str, *, fixed_parameters: dict[str, Any] = {}) -> pl.DataFrame:
         """