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

squirrels/_dashboards.py

@@ -0,0 +1,179 @@
+from typing import Type, TypeVar, Callable, Coroutine, Any
+from typing_extensions import Self
+from enum import Enum
+from dataclasses import dataclass
+from pydantic import BaseModel, Field, model_validator, ValidationError
+import matplotlib.figure as figure
+import os, time, io, abc, typing
+
+from ._arguments.run_time_args import DashboardArgs
+from ._py_module import PyModule
+from ._manifest import AnalyticsOutputConfig, AuthType, PermissionScope, ConfigurableOverride
+from ._exceptions import InvalidInputError, ConfigurationError, FileExecutionError
+from . import _constants as c, _utils as u
+
+
+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) -> typing.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) -> typing.Literal['html']:
+        return c.HTML
+    
+    def _repr_html_(self):
+        return self._content
+
+
+T = TypeVar('T', bound=Dashboard)
+
+
+class DashboardFormat(Enum):
+    PNG = "png"
+    HTML = "html"
+
+class DashboardDependencies(BaseModel):
+    name: str
+    dataset: str
+    fixed_parameters: dict[str, str] = Field(default_factory=dict)
+
+class DashboardConfig(AnalyticsOutputConfig):
+    format: DashboardFormat = Field(default=DashboardFormat.PNG)
+    depends_on: list[DashboardDependencies] = Field(default_factory=list)
+
+
+@dataclass
+class DashboardDefinition:
+    dashboard_name: str
+    filepath: str
+    config: DashboardConfig
+    project_path: str
+
+    @property
+    def dashboard_func(self) -> Callable[[DashboardArgs], Coroutine[Any, Any, Dashboard]]:
+        if not hasattr(self, '_dashboard_func'):
+            module = PyModule(self.filepath, self.project_path)
+            self._dashboard_func = module.get_func_or_class(c.MAIN_FUNC)
+        return self._dashboard_func
+
+    def get_dashboard_format(self) -> str:
+        return self.config.format.value
+    
+    async def get_dashboard(self, args: DashboardArgs, *, dashboard_type: Type[T] = Dashboard) -> T:
+        try:
+            dashboard = await self.dashboard_func(args)
+            assert isinstance(dashboard, dashboard_type), f"Function does not return expected Dashboard type: {dashboard_type}"
+        except (InvalidInputError, ConfigurationError, FileExecutionError) as e:
+            raise e
+        except Exception as e:
+            raise FileExecutionError(f'Failed to run "{c.MAIN_FUNC}" function for dashboard "{self.dashboard_name}"', e) from e
+
+        return dashboard
+
+
+class DashboardsIO:
+
+    @classmethod
+    def load_files(
+        cls, logger: u.Logger, project_path: str, auth_type: AuthType = AuthType.OPTIONAL, project_configurables: dict[str, Any] | None = None
+    ) -> dict[str, DashboardDefinition]:
+        start = time.time()
+        
+        if project_configurables is None:
+            project_configurables = {}
+        
+        default_scope = PermissionScope.PROTECTED if auth_type == AuthType.REQUIRED else PermissionScope.PUBLIC
+        
+        dashboards_by_name = {}
+        for dp, _, filenames in os.walk(u.Path(project_path, c.DASHBOARDS_FOLDER)):
+            for file in filenames:
+                filepath = os.path.join(dp, file)
+                file_stem, extension = os.path.splitext(file)
+                if not extension == '.py':
+                    continue
+                
+                # Check for corresponding .yml file
+                yml_path = os.path.join(dp, file_stem + '.yml')
+                config_dict = u.load_yaml_config(yml_path) if os.path.exists(yml_path) else {}
+                try:
+                    config = DashboardConfig(name=file_stem, project_configurables=project_configurables, **config_dict)
+                except ValidationError as e:
+                    raise ConfigurationError(f'Failed to process dashboard config "{file_stem}". ' + str(e)) from e
+                
+                if config.scope is None:
+                    config.scope = default_scope
+                
+                if auth_type == AuthType.REQUIRED and config.scope == PermissionScope.PUBLIC:
+                    raise ConfigurationError(f'Authentication is required, so dashboard "{file_stem}" cannot be public. Update the scope in "{yml_path}"')
+                
+                dashboards_by_name[file_stem] = DashboardDefinition(file_stem, filepath, config, project_path)
+                
+        logger.log_activity_time("loading files for dashboards", start)
+        return dashboards_by_name

squirrels/_data_sources.py

@@ -1,10 +1,16 @@
 from __future__ import annotations
 from dataclasses import dataclass
-import polars as pl, typing as t, abc
+from enum import Enum
+import polars as pl, typing as t, abc, re
 
 from . import _parameter_configs as pc, _parameter_options as po
 from ._exceptions import ConfigurationError
 
+class SourceEnum(Enum):
+    CONNECTION = "connection"
+    SEEDS = "seeds"
+    VDL = "vdl"
+
 
 @dataclass
 class DataSource(metaclass=abc.ABCMeta):
@@ -13,19 +19,19 @@ class DataSource(metaclass=abc.ABCMeta):
     """
     _table_or_query: str
     _id_col: str | None
-    _is_from_seeds: bool
+    _source: SourceEnum
     _user_group_col: str | None
     _parent_id_col: str | None
     _connection: str | None
 
     @abc.abstractmethod
     def __init__(
-        self, table_or_query: str, *, id_col: str | None = None, from_seeds: bool = False, user_group_col: str | None = None, 
-        parent_id_col: str | None = None, connection: str | None = None, **kwargs
+        self, table_or_query: str, *, id_col: str | None = None, source: SourceEnum = SourceEnum.CONNECTION, 
+        user_group_col: str | None = None, parent_id_col: str | None = None, connection: str | None = None, **kwargs
     ) -> None:
         self._table_or_query = table_or_query
         self._id_col = id_col
-        self._is_from_seeds = from_seeds
+        self._source = source
         self._user_group_col = user_group_col
         self._parent_id_col = parent_id_col
         self._connection = connection
@@ -40,7 +46,7 @@ class DataSource(metaclass=abc.ABCMeta):
         Returns:
             str: The converted select query
         """
-        if self._table_or_query.strip().lower().startswith('select '):
+        if re.match(r'select\s', self._table_or_query.strip(), re.IGNORECASE):
             query = self._table_or_query
         else:
             query = f'SELECT * FROM {self._table_or_query}'
@@ -100,12 +106,12 @@ class _SelectionDataSource(DataSource):
     @abc.abstractmethod
     def __init__(
         self, table_or_query: str, id_col: str, options_col: str, *, order_by_col: str | None = None, 
-        is_default_col: str | None = None, custom_cols: dict[str, str] = {}, from_seeds: bool = False, 
+        is_default_col: str | None = None, custom_cols: dict[str, str] = {}, source: SourceEnum = SourceEnum.CONNECTION, 
         user_group_col: str | None = None, parent_id_col: str | None = None, connection: str | None = None, 
         **kwargs
     ) -> None:
         super().__init__(
-            table_or_query, id_col=id_col, from_seeds=from_seeds, user_group_col=user_group_col, parent_id_col=parent_id_col, 
+            table_or_query, id_col=id_col, source=source, user_group_col=user_group_col, parent_id_col=parent_id_col, 
             connection=connection
         )
         self._options_col = options_col
@@ -151,7 +157,7 @@ class SelectDataSource(_SelectionDataSource):
 
     def __init__(
             self, table_or_query: str, id_col: str, options_col: str, *, order_by_col: str | None = None, 
-            is_default_col: str | None = None, custom_cols: dict[str, str] = {}, from_seeds: bool = False, 
+            is_default_col: str | None = None, custom_cols: dict[str, str] = {}, source: SourceEnum = SourceEnum.CONNECTION, 
             user_group_col: str | None = None, parent_id_col: str | None = None, connection: str | None = None, 
             **kwargs
         ) -> None:
@@ -165,19 +171,19 @@ class SelectDataSource(_SelectionDataSource):
             order_by_col: The column name to order the options by. Orders by the id_col instead if this is None
             is_default_col: The column name that indicates which options are the default
             custom_cols: Dictionary of attribute to column name for custom fields for the SelectParameterOption
-            from_seeds: Boolean for whether this datasource is created from seeds
+            source: The source to fetch data from. Must be "connection", "seeds", or "vdl". Defaults to "connection"
             user_group_col: The column name of the user group that the user is in for this option to be valid
             parent_id_col: The column name of the parent option id that must be selected for this option to be valid
             connection: Name of the connection to use defined in connections.py
         """
         super().__init__(
             table_or_query, id_col, options_col, order_by_col=order_by_col, is_default_col=is_default_col, custom_cols=custom_cols,
-            from_seeds=from_seeds, user_group_col=user_group_col, parent_id_col=parent_id_col, connection=connection
+            source=source, user_group_col=user_group_col, parent_id_col=parent_id_col, connection=connection
         )
 
     def _convert(self, ds_param: pc.DataSourceParameterConfig, df: pl.DataFrame) -> pc.SelectionParameterConfig:
         """
-        Method to convert the associated DataSourceParameter into a SingleSelectParameterConfig or MultiSelectParameterConfig
+        Method to convert the associated DataSourceParameterConfig into a SingleSelectParameterConfig or MultiSelectParameterConfig
 
         Arguments:
             ds_param: The parameter to convert
@@ -212,7 +218,7 @@ class DateDataSource(DataSource):
     def __init__(
         self, table_or_query: str, default_date_col: str, *, min_date_col: str | None = None, 
         max_date_col: str | None = None, date_format: str = '%Y-%m-%d', id_col: str | None = None, 
-        from_seeds: bool = False, user_group_col: str | None = None, parent_id_col: str | None = None, 
+        source: SourceEnum = SourceEnum.CONNECTION, user_group_col: str | None = None, parent_id_col: str | None = None, 
         connection: str | None = None, **kwargs
     ) -> None:
         """
@@ -223,13 +229,13 @@ class DateDataSource(DataSource):
             default_date_col: The column name of the default date
             date_format: The format of the default date(s). Defaults to '%Y-%m-%d'
             id_col: The column name of the id
-            from_seeds: Boolean for whether this datasource is created from seeds
+            source: The source to fetch data from. Must be "connection", "seeds", or "vdl". Defaults to "connection"
             user_group_col: The column name of the user group that the user is in for this option to be valid
             parent_id_col: The column name of the parent option id that the default date belongs to
             connection: Name of the connection to use defined in connections.py
         """
         super().__init__(
-            table_or_query, id_col=id_col, from_seeds=from_seeds, user_group_col=user_group_col, parent_id_col=parent_id_col, 
+            table_or_query, id_col=id_col, source=source, user_group_col=user_group_col, parent_id_col=parent_id_col, 
             connection=connection
         )
         self._default_date_col = default_date_col
@@ -239,7 +245,7 @@ class DateDataSource(DataSource):
 
     def _convert(self, ds_param: pc.DataSourceParameterConfig, df: pl.DataFrame) -> pc.DateParameterConfig:
         """
-        Method to convert the associated DataSourceParameter into a DateParameterConfig
+        Method to convert the associated DataSourceParameterConfig into a DateParameterConfig
 
         Arguments:
             ds_param: The parameter to convert
@@ -281,7 +287,7 @@ class DateRangeDataSource(DataSource):
 
     def __init__(
         self, table_or_query: str, default_start_date_col: str, default_end_date_col: str, *, date_format: str = '%Y-%m-%d',
-        min_date_col: str | None = None, max_date_col: str | None = None, id_col: str | None = None, from_seeds: bool = False, 
+        min_date_col: str | None = None, max_date_col: str | None = None, id_col: str | None = None, source: SourceEnum = SourceEnum.CONNECTION, 
         user_group_col: str | None = None, parent_id_col: str | None = None, connection: str | None = None, **kwargs
     ) -> None:
         """
@@ -293,13 +299,13 @@ class DateRangeDataSource(DataSource):
             default_end_date_col: The column name of the default end date
             date_format: The format of the default date(s). Defaults to '%Y-%m-%d'
             id_col: The column name of the id
-            from_seeds: Boolean for whether this datasource is created from seeds
+            source: The source to fetch data from. Must be "connection", "seeds", or "vdl". Defaults to "connection"
             user_group_col: The column name of the user group that the user is in for this option to be valid
             parent_id_col: The column name of the parent option id that the default date belongs to
             connection: Name of the connection to use defined in connections.py
         """
         super().__init__(
-            table_or_query, id_col=id_col, from_seeds=from_seeds, user_group_col=user_group_col, parent_id_col=parent_id_col, 
+            table_or_query, id_col=id_col, source=source, user_group_col=user_group_col, parent_id_col=parent_id_col, 
             connection=connection
         )
         self._default_start_date_col = default_start_date_col
@@ -310,7 +316,7 @@ class DateRangeDataSource(DataSource):
 
     def _convert(self, ds_param: pc.DataSourceParameterConfig, df: pl.DataFrame) -> pc.DateRangeParameterConfig:
         """
-        Method to convert the associated DataSourceParameter into a DateRangeParameterConfig
+        Method to convert the associated DataSourceParameterConfig into a DateRangeParameterConfig
 
         Arguments:
             ds_param: The parameter to convert
@@ -354,11 +360,11 @@ class _NumericDataSource(DataSource):
     @abc.abstractmethod
     def __init__(
         self, table_or_query: str, min_value_col: str, max_value_col: str, *, increment_col: str | None = None, 
-        id_col: str | None = None, from_seeds: bool = False, user_group_col: str | None = None, 
+        id_col: str | None = None, source: SourceEnum = SourceEnum.CONNECTION, user_group_col: str | None = None, 
         parent_id_col: str | None = None, connection: str | None = None, **kwargs
     ) -> None:
         super().__init__(
-            table_or_query, id_col=id_col, from_seeds=from_seeds, user_group_col=user_group_col, parent_id_col=parent_id_col, 
+            table_or_query, id_col=id_col, source=source, user_group_col=user_group_col, parent_id_col=parent_id_col, 
             connection=connection
         )
         self._min_value_col = min_value_col
@@ -375,7 +381,7 @@ class NumberDataSource(_NumericDataSource):
 
     def __init__(
         self, table_or_query: str, min_value_col: str, max_value_col: str, *, increment_col: str | None = None,
-        default_value_col: str | None = None, id_col: str | None = None, from_seeds: bool = False, 
+        default_value_col: str | None = None, id_col: str | None = None, source: SourceEnum = SourceEnum.CONNECTION, 
         user_group_col: str | None = None, parent_id_col: str | None = None, connection: str | None = None, **kwargs
     ) -> None:
         """
@@ -388,20 +394,20 @@ class NumberDataSource(_NumericDataSource):
             increment_col: The column name of the increment value. Defaults to column of 1's if None
             default_value_col: The column name of the default value. Defaults to min_value_col if None
             id_col: The column name of the id
-            from_seeds: Boolean for whether this datasource is created from seeds
+            source: The source to fetch data from. Must be "connection", "seeds", or "vdl". Defaults to "connection"
             user_group_col: The column name of the user group that the user is in for this option to be valid
             parent_id_col: The column name of the parent option id that the default value belongs to
             connection: Name of the connection to use defined in connections.py
         """
         super().__init__(
-            table_or_query, min_value_col, max_value_col, increment_col=increment_col, id_col=id_col, from_seeds=from_seeds,
+            table_or_query, min_value_col, max_value_col, increment_col=increment_col, id_col=id_col, source=source,
             user_group_col=user_group_col, parent_id_col=parent_id_col, connection=connection
         )
         self._default_value_col = default_value_col
 
     def _convert(self, ds_param: pc.DataSourceParameterConfig, df: pl.DataFrame) -> pc.NumberParameterConfig:
         """
-        Method to convert the associated DataSourceParameter into a NumberParameterConfig
+        Method to convert the associated DataSourceParameterConfig into a NumberParameterConfig
 
         Arguments:
             ds_param: The parameter to convert
@@ -443,7 +449,7 @@ class NumberRangeDataSource(_NumericDataSource):
     def __init__(
         self, table_or_query: str, min_value_col: str, max_value_col: str, *, increment_col: str | None = None,
         default_lower_value_col: str | None = None, default_upper_value_col: str | None = None, id_col: str | None = None, 
-        from_seeds: bool = False, user_group_col: str | None = None, parent_id_col: str | None = None, 
+        source: SourceEnum = SourceEnum.CONNECTION, user_group_col: str | None = None, parent_id_col: str | None = None, 
         connection: str | None = None, **kwargs
     ) -> None:
         """
@@ -457,13 +463,13 @@ class NumberRangeDataSource(_NumericDataSource):
             default_lower_value_col: The column name of the default lower value. Defaults to min_value_col if None
             default_upper_value_col: The column name of the default upper value. Defaults to max_value_col if None
             id_col: The column name of the id
-            from_seeds: Boolean for whether this datasource is created from seeds
+            source: The source to fetch data from. Must be "connection", "seeds", or "vdl". Defaults to "connection"
             user_group_col: The column name of the user group that the user is in for this option to be valid
             parent_id_col: The column name of the parent option id that the default value belongs to
             connection: Name of the connection to use defined in connections.py
         """
         super().__init__(
-            table_or_query, min_value_col, max_value_col, increment_col=increment_col, id_col=id_col, from_seeds=from_seeds, 
+            table_or_query, min_value_col, max_value_col, increment_col=increment_col, id_col=id_col, source=source, 
             user_group_col=user_group_col, parent_id_col=parent_id_col, connection=connection
         )
         self._default_lower_value_col = default_lower_value_col
@@ -471,7 +477,7 @@ class NumberRangeDataSource(_NumericDataSource):
 
     def _convert(self, ds_param: pc.DataSourceParameterConfig, df: pl.DataFrame) -> pc.NumberRangeParameterConfig:
         """
-        Method to convert the associated DataSourceParameter into a NumberRangeParameterConfig
+        Method to convert the associated DataSourceParameterConfig into a NumberRangeParameterConfig
 
         Arguments:
             ds_param: The parameter to convert
@@ -511,7 +517,7 @@ class TextDataSource(DataSource):
     _default_text_col: str
 
     def __init__(
-        self, table_or_query: str, default_text_col: str, *, id_col: str | None = None, from_seeds: bool = False, 
+        self, table_or_query: str, default_text_col: str, *, id_col: str | None = None, source: SourceEnum = SourceEnum.CONNECTION, 
         user_group_col: str | None = None, parent_id_col: str | None = None, connection: str | None = None,
         **kwargs
     ) -> None:
@@ -522,20 +528,20 @@ class TextDataSource(DataSource):
             table_or_query: Either the name of the table to use, or a query to run
             default_text_col: The column name of the default text
             id_col: The column name of the id
-            from_seeds: Boolean for whether this datasource is created from seeds
+            source: The source to fetch data from. Must be "connection", "seeds", or "vdl". Defaults to "connection"
             user_group_col: The column name of the user group that the user is in for this option to be valid
             parent_id_col: The column name of the parent option id that the default date belongs to
             connection: Name of the connection to use defined in connections.py
         """
         super().__init__(
-            table_or_query, id_col=id_col, from_seeds=from_seeds, user_group_col=user_group_col, parent_id_col=parent_id_col, 
+            table_or_query, id_col=id_col, source=source, user_group_col=user_group_col, parent_id_col=parent_id_col, 
             connection=connection
         )
         self._default_text_col = default_text_col
 
     def _convert(self, ds_param: pc.DataSourceParameterConfig, df: pl.DataFrame) -> pc.TextParameterConfig:
         """
-        Method to convert the associated DataSourceParameter into a TextParameterConfig
+        Method to convert the associated DataSourceParameterConfig into a TextParameterConfig
 
         Arguments:
             ds_param: The parameter to convert

squirrels/_dataset_types.py

@@ -32,29 +32,34 @@ class DatasetMetadata:
         return self._json_repr
 
 
+@dataclass(frozen=True)
+class DatasetResultFormat:
+    orientation: Literal["records", "rows", "columns"]
+    offset: int
+    limit: int | None
+
+
 @dataclass
 class DatasetResult(DatasetMetadata):
     df: pl.DataFrame
-    to_json: Callable[[str, tuple[str, ...], int, int], dict] = field(init=False)
+    to_json: Callable[[DatasetResultFormat], dict] = field(init=False)
 
     def __post_init__(self):
         self.to_json = lru_cache()(self._to_json)
     
-    def _to_json(self, orientation: Literal["records", "rows", "columns"], select: tuple[str, ...], limit: int, offset: int) -> dict:
+    def _to_json(self, result_format: DatasetResultFormat) -> dict:
         df = self.df.lazy()
-        if offset > 0:
-            df = df.filter(pl.col("_row_num") > offset)
-        if limit > 0:
-            df = df.limit(limit)
-        if select:
-            df = df.select(select)
+        if result_format.offset > 0:
+            df = df.filter(pl.col("_row_num") > result_format.offset)
+        if result_format.limit is not None:
+            df = df.limit(result_format.limit)
         df = df.collect()
         
-        if orientation == "columns":
+        if result_format.orientation == "columns":
             data = df.to_dict(as_series=False)
         else:
             data = df.to_dicts()
-            if orientation == "rows":
+            if result_format.orientation == "rows":
                 data = [[row[col] for col in df.columns] for row in data]
 
         column_details_by_name = {col.name: col for col in self.target_model_config.columns}
@@ -80,7 +85,7 @@ class DatasetResult(DatasetMetadata):
             "total_num_rows": self.df.select(pl.len()).item(),
             "data_details": {
                 "num_rows": df.select(pl.len()).item(),
-                "orientation": orientation
+                "orientation": result_format.orientation
             },
             "data": data
         }