Fix type hint (#194)

Co-authored-by: sfc-gh-jsummer <[email protected]>

app_utils/shared_utils.py

@@ -7,7 +7,7 @@
 from datetime import datetime
 from enum import Enum
 from io import StringIO
-from typing import Any, Optional, List
+from typing import Any, Optional, List, Union
 
 import pandas as pd
 import streamlit as st
@@ -200,7 +200,7 @@ def get_available_stages(schema: str) -> List[str]:
     return fetch_stages_in_schema(get_snowflake_connection(), schema)
 
 
-def stage_selector_container() -> None | List[str]:
+def stage_selector_container() -> Optional[List[str]]:
     """
     Common component that encapsulates db/schema/stage selection for the admin app.
     When a db/schema/stage is selected, it is saved to the session state for reading elsewhere.
@@ -1247,23 +1247,23 @@ class AppMetadata:
     """
 
     @property
-    def user(self) -> str | None:
+    def user(self) -> Optional[str]:
         return os.getenv("SNOWFLAKE_USER")
 
     @property
-    def stage(self) -> str | None:
+    def stage(self) -> Optional[str]:
         if stage_exists():
             stage = st.session_state.snowflake_stage
             return f"{stage.stage_database}.{stage.stage_schema}.{stage.stage_name}"
         return None
 
     @property
-    def model(self) -> str | None:
+    def model(self) -> Optional[str]:
         if semantic_model_exists():
             return st.session_state.semantic_model.name  # type: ignore
         return None
 
-    def to_dict(self) -> dict[str, str | None]:
+    def to_dict(self) -> dict[str, Union[str,None]]:
         return {
             "User": self.user,
             "Stage": self.stage,

partner/cortex.py

@@ -18,11 +18,11 @@ def __init__(self, data: dict[str, Any]):
 
         self.data: dict[str, Any] = data
         self.name: str = data["name"]
-        self.synonyms: None | list[str] = data.get("synonyms", None)
+        self.synonyms: Optional[list[str]] = data.get("synonyms", None)
         self.data_type: str = data.get("data_type", "TEXT")
         self.expr: str = data["expr"]
-        self.description: None | str = data.get("description", None)
-        self.sample_values: None | list[str] = data.get("sample_values", None)
+        self.description: Optional[str] = data.get("description", None)
+        self.sample_values: Optional[list[str]] = data.get("sample_values", None)
         self.unique: bool = data.get("unique", False)
 
     def get_name(self) -> str:
@@ -90,13 +90,13 @@ class CortexSemanticTable:
     def __init__(self, data: dict[str, Any]):
         self.data: dict[str, Any] = data
         self.name: str = data["name"]
-        self.description: None | str = data["description"]
+        self.description: Optional[str] = data["description"]
         self.base_table_db: str = data["base_table"]["database"]
         self.base_table_schema: str = data["base_table"]["schema"]
         self.base_table_table: str = data["base_table"]["table"]
-        self.dimensions: None | list[dict[str, Any]] = data["dimensions"]
-        self.time_dimensions: None | list[dict[str, Any]] = data["time_dimensions"]
-        self.measures: None | list[dict[str, Any]] = data["measures"]
+        self.dimensions: Optional[list[dict[str, Any]]] = data["dimensions"]
+        self.time_dimensions: Optional[list[dict[str, Any]]] = data["time_dimensions"]
+        self.measures: Optional[list[dict[str, Any]]] = data["measures"]
 
     def get_data(self) -> dict[str, Any]:
         return self.data

partner/dbt.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 import pandas as pd
 import streamlit as st
@@ -80,7 +80,7 @@ def upload_dbt_semantic() -> None:
             key="dbt_files",
         )
     if uploaded_files:
-        partner_semantic: list[None | DBTSemanticModel] = []
+        partner_semantic: list[Union[None,DBTSemanticModel]] = []
         for file in uploaded_files:
             partner_semantic.extend(read_dbt_yaml(file))  # type: ignore
 
@@ -111,7 +111,7 @@ def __init__(self, entity: dict[str, Any]):
         self.name: str = entity["name"]
         self.type: str = entity.get("type", None)
         self.expr: str = entity.get("expr", self.name)
-        self.description: None | str = entity.get("description", None)
+        self.description: Optional[str] = entity.get("description", None)
         self.cortex_map = {
             "name": self.name,
             "description": self.description,
@@ -153,7 +153,7 @@ class DBTMeasure(DBTEntity):
 
     def __init__(self, entity: dict[str, Any]):
         super().__init__(entity)
-        self.agg: None | str = entity.get("agg", None)
+        self.agg: Optional[str] = entity.get("agg", None)
         self.cortex_map = {
             "name": self.name,
             "description": self.description,
@@ -195,10 +195,10 @@ class DBTSemanticModel:
     def __init__(self, data: dict[str, Any]):
         self.data: dict[str, Any] = data
         self.name: str = data["name"]
-        self.description: None | str = data.get("description", None)
-        self.entities: None | list[dict[str, Any]] = data["entities"]
-        self.dimensions: None | list[dict[str, Any]] = data["dimensions"]
-        self.measures: None | list[dict[str, Any]] = data["measures"]
+        self.description: Optional[str] = data.get("description", None)
+        self.entities: Optional[list[dict[str, Any]]] = data["entities"]
+        self.dimensions: Optional[list[dict[str, Any]]] = data["dimensions"]
+        self.measures: Optional[list[dict[str, Any]]] = data["measures"]
 
     def get_data(self) -> dict[str, Any]:
         return self.data

partner/looker.py

@@ -1,5 +1,5 @@
 import os
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 import pandas as pd
 import streamlit as st
@@ -523,7 +523,7 @@ def render_looker_explore_as_table(
     target_lag: Optional[int] = 20,
     target_lag_unit: Optional[str] = "minutes",
     warehouse: Optional[str] = None,
-) -> None | dict[str, dict[str, str]]:
+) -> Union[None,dict[str, dict[str, str]]]:
     """
     Creates materialized table corresponding to Looker Explore.
     Args:

partner/partner_utils.py

@@ -113,7 +113,7 @@ def __init__(self, row_data: pd.Series) -> None:  # type: ignore
             else {}
         )
 
-    def render_row(self) -> None | dict[str, Any]:  # type: ignore
+    def render_row(self) -> Union[None, dict[str, Any]]:  # type: ignore
         toggle_options = ["merged", "cortex", "partner", "remove"]
         metadata = {}
 

semantic_model_generator/snowflake_utils/snowflake_connector.py

@@ -129,7 +129,7 @@ def _get_column_comment(
 def get_table_primary_keys(
     conn: SnowflakeConnection,
     table_fqn: str,
-) -> list[str] | None:
+) -> Optional[list[str]]:
     query = f"show primary keys in table {table_fqn};"
     cursor = conn.cursor()
     cursor.execute(query)