squirrels 0.3.2__py3-none-any.whl → 0.4.0__py3-none-any.whl

squirrels/arguments/run_time_args.py

@@ -1,4 +1,4 @@
-from typing import Union, Callable, Optional, Any
+from typing import Iterable, Callable, Any, Coroutine
 from dataclasses import dataclass
 from sqlalchemy import Engine
 import pandas as pd
@@ -6,29 +6,58 @@ import pandas as pd
 from .init_time_args import ConnectionsArgs, ParametersArgs
 from ..user_base import User
 from ..parameters import Parameter, TextValue
-from .._connection_set import ConnectionSetIO
-from .. import _utils as u
+from .. import _utils as _u
 
 
 @dataclass
 class AuthArgs(ConnectionsArgs):
-    connections: dict[str, Engine]
+    _connections: dict[str, Engine]
     username: str
     password: str
 
+    @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()
+
 
 @dataclass
 class ContextArgs(ParametersArgs):
-    user: User
-    prms: dict[str, Parameter]
-    traits: dict[str, Any]
+    user: User | None
+    _prms: dict[str, Parameter]
+    _traits: dict[str, Any]
     _placeholders: dict[str, Any]
 
-    def set_placeholder(self, placeholder: str, value: Union[TextValue, Any]) -> str:
+    @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(self) -> dict[str, Any]:
+        """
+        A dictionary of placeholder name to placeholder value
+        """
+        return self._placeholders.copy()
+
+    def set_placeholder(self, placeholder: str, value: TextValue | Any) -> str:
         """
         Method to set a placeholder value.
 
-        Parameters:
+        Arguments:
             placeholder: A string for the name of the placeholder
             value: The value of the placeholder. Can be of any type
         """
@@ -42,7 +71,7 @@ class ContextArgs(ParametersArgs):
         Method to check whether a given parameter exists and is enabled (i.e., not hidden based on other parameter selections) for the current 
         dataset at runtime.
 
-        Parameters:
+        Arguments:
             param_name: A string for the name of the parameter
         
         Returns:
@@ -53,29 +82,46 @@ class ContextArgs(ParametersArgs):
 
 @dataclass
 class ModelDepsArgs(ContextArgs):
-    ctx: dict[str, Any]
+    _ctx: dict[str, Any]
+
+    @property
+    def ctx(self) -> dict[str, Any]:
+        """
+        Dictionary of context variables
+        """
+        return self._ctx.copy()
 
 
 @dataclass
 class ModelArgs(ModelDepsArgs):
     connection_name: str
     _connections: dict[str, Engine]
-    _dependencies: set[str]
+    _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]:
-        return self._dependencies.copy()
+        """
+        The set of dependent data model names
+        """
+        return set(self._dependencies)
     
     def is_placeholder(self, placeholder: str) -> bool:
         """
         Checks whether a name is a valid placeholder
 
-        Parameters:
+        Arguments:
             placeholder: A string for the name of the placeholder
         
         Returns:
@@ -83,13 +129,13 @@ class ModelArgs(ModelDepsArgs):
         """
         return placeholder in self._placeholders
     
-    def get_placeholder_value(self, placeholder: str) -> Optional[Any]:
+    def get_placeholder_value(self, placeholder: str) -> Any | None:
         """
         Gets the value of a placeholder.
 
         USE WITH CAUTION. Do not use the return value directly in a SQL query since that could be prone to SQL injection
 
-        Parameters:
+        Arguments:
             placeholder: A string for the name of the placeholder
         
         Returns:
@@ -104,7 +150,7 @@ class ModelArgs(ModelDepsArgs):
         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.
 
-        Parameters:
+        Arguments:
             model: The model name
         
         Returns:
@@ -112,25 +158,24 @@ class ModelArgs(ModelDepsArgs):
         """
         return self._ref(model)
 
-    def run_external_sql(self, sql_query: str, *, connection_name: str = None, **kwargs) -> pd.DataFrame:
+    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
 
-        Parameters:
+        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
         """
-        connection_name = self.connection_name if connection_name is None else connection_name
-        return ConnectionSetIO.obj.run_sql_query_from_conn_name(sql_query, connection_name, self._placeholders)
+        return self._run_external_sql(sql_query, connection_name)
 
-    def run_sql_on_dataframes(self, sql_query: str, *, dataframes: dict[str, pd.DataFrame] = None, **kwargs) -> pd.DataFrame:
+    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)
 
-        Parameters:
+        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
         
@@ -140,4 +185,24 @@ class ModelArgs(ModelDepsArgs):
         if dataframes is None:
             dataframes = {x: self.ref(x) for x in self._dependencies}
 
-        return u.run_sql_on_dataframes(sql_query, dataframes)
+        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:
+        """
+        Get dataset as DataFrame given dataset name. Can use this to access protected/private datasets regardless of user authenticated to the dashboard.
+
+        The parameters used for the dataset include the parameter selections coming from the REST API and the fixed_parameters argument. The fixed_parameters takes precedence.
+
+        Arguments:
+            name: A string for the dataset name
+            fixed_parameters: Parameters to set for this dataset (in addition to the ones set through real-time selections)
+        
+        Returns:
+            A DataFrame for the result of the dataset
+        """
+        return await self._get_dataset(name, fixed_parameters)

squirrels/dashboards.py

@@ -0,0 +1,82 @@
+import matplotlib.figure as _figure, io as _io, abc as _abc, typing as _t
+
+from . import _constants as _c
+
+
+class Dashboard(metaclass=_abc.ABCMeta):
+    """
+    Abstract parent class for all Dashboard classes.
+    """
+    
+    @property
+    @_abc.abstractmethod
+    def _content(self) -> bytes | str:
+        pass
+    
+    @property
+    @_abc.abstractmethod
+    def _format(self) -> str:
+        pass
+
+
+class PngDashboard(Dashboard):
+    """
+    Instantiate a Dashboard in PNG format from a matplotlib figure or bytes
+    """
+    
+    def __init__(self, content: _figure.Figure | _io.BytesIO | bytes) -> None:
+        """
+        Constructor for PngDashboard
+
+        Arguments:
+            content: The content of the dashboard as a matplotlib.figure.Figure or bytes
+        """
+        if isinstance(content, _figure.Figure):
+            buffer = _io.BytesIO()
+            content.savefig(buffer, format=_c.PNG)
+            content = buffer.getvalue()
+        
+        if isinstance(content, _io.BytesIO):
+            content = content.getvalue()
+        
+        self.__content = content
+
+    @property
+    def _content(self) -> bytes:
+        return self.__content
+    
+    @property
+    def _format(self) -> _t.Literal['png']:
+        return _c.PNG
+    
+    def _repr_png_(self):
+        return self._content
+    
+
+class HtmlDashboard(Dashboard):
+    """
+    Instantiate a Dashboard from an HTML string
+    """
+
+    def __init__(self, content: _io.StringIO | str) -> None:
+        """
+        Constructor for HtmlDashboard
+
+        Arguments:
+            content: The content of the dashboard as HTML string
+        """
+        if isinstance(content, _io.StringIO):
+            content = content.getvalue()
+        
+        self.__content = content
+
+    @property
+    def _content(self) -> str:
+        return self.__content
+    
+    @property
+    def _format(self) -> _t.Literal['html']:
+        return _c.HTML
+    
+    def _repr_html_(self):
+        return self._content