squirrels 0.4.1__py3-none-any.whl → 0.5.0rc0__py3-none-any.whl

squirrels/arguments/run_time_args.py

@@ -1,36 +1,46 @@
-from typing import Iterable, Callable, Any, Coroutine
-from dataclasses import dataclass
-from sqlalchemy import Engine
-import pandas as pd
+from typing import Callable, Any, Coroutine
+import polars as pl
 
-from .init_time_args import ConnectionsArgs, ParametersArgs
-from ..user_base import User
+from .init_time_args import _WithConnectionDictArgs, ParametersArgs, BuildModelArgs, ConnectionsArgs
+from .._auth import BaseUser
 from ..parameters import Parameter, TextValue
-from .. import _utils as _u
 
 
-@dataclass
-class AuthArgs(ConnectionsArgs):
-    _connections: dict[str, Engine]
-    username: str
-    password: str
+class AuthLoginArgs(_WithConnectionDictArgs):
 
-    @property
-    def connections(self) -> dict[str, Engine]:
-        """
-        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()
+    def __init__(
+        self, conn_args: ConnectionsArgs, _connections: dict[str, Any], 
+        username: str, 
+        password: str
+    ):
+        super().__init__(conn_args.project_path, conn_args.proj_vars, conn_args.env_vars, _connections)
+        self.username = username
+        self.password = password
+
+
+class AuthTokenArgs(_WithConnectionDictArgs):
+
+    def __init__(
+        self, conn_args: ConnectionsArgs, _connections: dict[str, Any], 
+        token: str
+    ):
+        super().__init__(conn_args.project_path, conn_args.proj_vars, conn_args.env_vars, _connections)
+        self.token = token
 
 
-@dataclass
 class ContextArgs(ParametersArgs):
-    user: User | None
-    _prms: dict[str, Parameter]
-    _traits: dict[str, Any]
-    _placeholders: dict[str, Any]
+
+    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]:
@@ -80,9 +90,25 @@ class ContextArgs(ParametersArgs):
         return (param_name in self.prms and self.prms[param_name].is_enabled())
 
 
-@dataclass
-class ModelDepsArgs(ContextArgs):
-    _ctx: dict[str, Any]
+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
+        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]:
@@ -90,32 +116,6 @@ class ModelDepsArgs(ContextArgs):
         Dictionary of context variables
         """
         return self._ctx.copy()
-
-
-@dataclass
-class ModelArgs(ModelDepsArgs):
-    connection_name: str
-    _connections: dict[str, Engine]
-    _dependencies: Iterable[str]
-    _ref: Callable[[str], pd.DataFrame]
-    _run_external_sql: Callable[[str, str | None], pd.DataFrame]
-    _use_duckdb: bool
-
-    @property
-    def connections(self) -> dict[str, Engine]:
-        """
-        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()
-
-    @property
-    def dependencies(self) -> set[str]:
-        """
-        The set of dependent data model names
-        """
-        return set(self._dependencies)
     
     def is_placeholder(self, placeholder: str) -> bool:
         """
@@ -142,57 +142,18 @@ class ModelArgs(ModelDepsArgs):
             An type for the value of the placeholder
         """
         return self._placeholders.get(placeholder)
-    
-    def ref(self, model: str) -> pd.DataFrame:
-        """
-        Returns the result (as pandas DataFrame) of a dependent model (predefined in "dependencies" function)
 
-        Note: This is different behaviour than the "ref" function for SQL models, which figures out the dependent models for you, 
-        and returns a string for the table/view name in SQLite instead of a pandas DataFrame.
 
-        Arguments:
-            model: The model name
-        
-        Returns:
-            A pandas DataFrame
-        """
-        return self._ref(model)
-
-    def run_external_sql(self, sql_query: str, *, connection_name: str | None = None, **kwargs) -> pd.DataFrame:
-        """
-        Runs a SQL query against an external database, with option to specify the connection name. Placeholder values are provided automatically
-
-        Arguments:
-            sql_query: The SQL query. Can be parameterized with placeholders
-            connection_name: The connection name for the database. If None, uses the one configured for the model
-        
-        Returns:
-            The query result as a pandas DataFrame
-        """
-        return self._run_external_sql(sql_query, connection_name)
-
-    def run_sql_on_dataframes(self, sql_query: str, *, dataframes: dict[str, pd.DataFrame] | None = None, **kwargs) -> pd.DataFrame:
-        """
-        Uses a dictionary of dataframes to execute a SQL query in an embedded in-memory database (sqlite or duckdb based on setting)
-
-        Arguments:
-            sql_query: The SQL query to run
-            dataframes: A dictionary of table names to their pandas Dataframe. If None, uses results of dependent models
-        
-        Returns:
-            The result as a pandas Dataframe from running the query
-        """
-        if dataframes is None:
-            dataframes = {x: self.ref(x) for x in self._dependencies}
-
-        return _u.run_sql_on_dataframes(sql_query, dataframes, self._use_duckdb)
-
-
-@dataclass
 class DashboardArgs(ParametersArgs):
-    _get_dataset: Callable[[str, dict[str, Any]], Coroutine[Any, Any, pd.DataFrame]]
 
-    async def dataset(self, name: str, *, fixed_parameters: dict[str, Any] = {}) -> pd.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
+
+    async def dataset(self, name: str, *, fixed_parameters: dict[str, Any] = {}) -> pl.DataFrame:
         """
         Get dataset as DataFrame given dataset name. Can use this to access protected/private datasets regardless of user authenticated to the dashboard.
 

squirrels/dashboards.py

@@ -1,6 +1,6 @@
 import matplotlib.figure as _figure, io as _io, abc as _abc, typing as _t
 
-from . import _constants as _c
+from . import _constants as c
 
 
 class Dashboard(metaclass=_abc.ABCMeta):
@@ -33,7 +33,7 @@ class PngDashboard(Dashboard):
         """
         if isinstance(content, _figure.Figure):
             buffer = _io.BytesIO()
-            content.savefig(buffer, format=_c.PNG)
+            content.savefig(buffer, format=c.PNG)
             content = buffer.getvalue()
         
         if isinstance(content, _io.BytesIO):
@@ -47,7 +47,7 @@ class PngDashboard(Dashboard):
     
     @property
     def _format(self) -> _t.Literal['png']:
-        return _c.PNG
+        return c.PNG
     
     def _repr_png_(self):
         return self._content
@@ -76,7 +76,7 @@ class HtmlDashboard(Dashboard):
     
     @property
     def _format(self) -> _t.Literal['html']:
-        return _c.HTML
+        return c.HTML
     
     def _repr_html_(self):
         return self._content