From 0aae2245149a82342d59fd899c386e45f8b93099 Mon Sep 17 00:00:00 2001 From: John Bodley Date: Sat, 4 Apr 2020 16:36:22 -0700 Subject: [PATCH] [mypy] Enforcing typing for superset.examples --- setup.cfg | 2 +- superset/examples/bart_lines.py | 2 +- superset/examples/birth_names.py | 10 +++++--- superset/examples/countries.py | 9 ++----- superset/examples/country_map.py | 2 +- superset/examples/css_templates.py | 2 +- superset/examples/deck.py | 2 +- superset/examples/energy.py | 2 +- superset/examples/flights.py | 2 +- superset/examples/helpers.py | 12 +++++---- superset/examples/long_lat.py | 2 +- superset/examples/misc_dashboard.py | 2 +- superset/examples/multi_line.py | 2 +- superset/examples/multiformat_time_series.py | 23 +++++++++-------- superset/examples/paris.py | 2 +- superset/examples/random_time_series.py | 4 ++- superset/examples/sf_population_polygons.py | 4 ++- superset/examples/tabbed_dashboard.py | 2 +- superset/examples/unicode_test_data.py | 2 +- superset/examples/world_bank.py | 6 ++--- superset/utils/core.py | 26 +++++++++++++++++--- superset/viz.py | 4 +-- 22 files changed, 75 insertions(+), 49 deletions(-) diff --git a/setup.cfg b/setup.cfg index 9fd8ef228f3c8..33e7223f1ddc5 100644 --- a/setup.cfg +++ b/setup.cfg @@ -53,7 +53,7 @@ order_by_type = false ignore_missing_imports = true no_implicit_optional = true -[mypy-superset.bin.*,superset.charts.*,superset.commands.*,superset.common.*,superset.dao.*,superset.db_engine_specs.*,superset.db_engines.*] +[mypy-superset.bin.*,superset.charts.*,superset.commands.*,superset.common.*,superset.dao.*,superset.db_engine_specs.*,superset.db_engines.*,superset.examples.*] check_untyped_defs = true disallow_untyped_calls = true disallow_untyped_defs = true diff --git a/superset/examples/bart_lines.py b/superset/examples/bart_lines.py index 0181d1ea92d7b..60b84b5e90436 100644 --- a/superset/examples/bart_lines.py +++ b/superset/examples/bart_lines.py @@ -26,7 +26,7 @@ from .helpers import get_example_data, TBL -def load_bart_lines(only_metadata=False, force=False): +def load_bart_lines(only_metadata: bool = False, force: bool = False) -> None: tbl_name = "bart_lines" database = get_example_database() table_exists = database.has_table_by_name(tbl_name) diff --git a/superset/examples/birth_names.py b/superset/examples/birth_names.py index e20c235dd026e..ade420b33db71 100644 --- a/superset/examples/birth_names.py +++ b/superset/examples/birth_names.py @@ -16,6 +16,7 @@ # under the License. import json import textwrap +from typing import Dict, Union import pandas as pd from sqlalchemy import DateTime, String @@ -23,6 +24,7 @@ from superset import db, security_manager from superset.connectors.sqla.models import SqlMetric, TableColumn +from superset.models.core import Database from superset.models.dashboard import Dashboard from superset.models.slice import Slice from superset.utils.core import get_example_database @@ -38,7 +40,9 @@ ) -def gen_filter(subject, comparator, operator="=="): +def gen_filter( + subject: str, comparator: str, operator: str = "==" +) -> Dict[str, Union[bool, str]]: return { "clause": "WHERE", "comparator": comparator, @@ -49,7 +53,7 @@ def gen_filter(subject, comparator, operator="=="): } -def load_data(tbl_name, database): +def load_data(tbl_name: str, database: Database) -> None: pdf = pd.read_json(get_example_data("birth_names.json.gz")) pdf.ds = pd.to_datetime(pdf.ds, unit="ms") pdf.to_sql( @@ -69,7 +73,7 @@ def load_data(tbl_name, database): print("-" * 80) -def load_birth_names(only_metadata=False, force=False): +def load_birth_names(only_metadata: bool = False, force: bool = False) -> None: """Loading birth name dataset from a zip file in the repo""" # pylint: disable=too-many-locals tbl_name = "birth_names" diff --git a/superset/examples/countries.py b/superset/examples/countries.py index 2bc352901d3f5..97238d8c7737e 100644 --- a/superset/examples/countries.py +++ b/superset/examples/countries.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. """This module contains data related to countries and is used for geo mapping""" -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional countries: List[Dict[str, Any]] = [ { @@ -2498,13 +2498,8 @@ all_lookups[lookup][country[lookup].lower()] = country -def get(field, symbol): +def get(field: str, symbol: str) -> Optional[Dict[str, Any]]: """ Get country data based on a standard code and a symbol - - >>> get('cioc', 'CUB')['name'] - "Cuba" - >>> get('cca2', 'CA')['name'] - "Canada" """ return all_lookups[field].get(symbol.lower()) diff --git a/superset/examples/country_map.py b/superset/examples/country_map.py index d77eb955a7c72..e9966ff305a4b 100644 --- a/superset/examples/country_map.py +++ b/superset/examples/country_map.py @@ -34,7 +34,7 @@ ) -def load_country_map_data(only_metadata=False, force=False): +def load_country_map_data(only_metadata: bool = False, force: bool = False) -> None: """Loading data for map with country map""" tbl_name = "birth_france_by_region" database = utils.get_example_database() diff --git a/superset/examples/css_templates.py b/superset/examples/css_templates.py index be677055ec402..4f3f355895ef9 100644 --- a/superset/examples/css_templates.py +++ b/superset/examples/css_templates.py @@ -20,7 +20,7 @@ from superset.models.core import CssTemplate -def load_css_templates(): +def load_css_templates() -> None: """Loads 2 css templates to demonstrate the feature""" print("Creating default CSS templates") diff --git a/superset/examples/deck.py b/superset/examples/deck.py index bce4df3be5e88..b9b85ee1faab7 100644 --- a/superset/examples/deck.py +++ b/superset/examples/deck.py @@ -167,7 +167,7 @@ }""" -def load_deck_dash(): +def load_deck_dash() -> None: print("Loading deck.gl dashboard") slices = [] tbl = db.session.query(TBL).filter_by(table_name="long_lat").first() diff --git a/superset/examples/energy.py b/superset/examples/energy.py index c3c33bd37f433..b27c45ebb523c 100644 --- a/superset/examples/energy.py +++ b/superset/examples/energy.py @@ -29,7 +29,7 @@ from .helpers import get_example_data, merge_slice, misc_dash_slices, TBL -def load_energy(only_metadata=False, force=False): +def load_energy(only_metadata: bool = False, force: bool = False) -> None: """Loads an energy related dataset to use with sankey and graphs""" tbl_name = "energy_usage" database = utils.get_example_database() diff --git a/superset/examples/flights.py b/superset/examples/flights.py index e4db4ca0c429d..16f88491dd340 100644 --- a/superset/examples/flights.py +++ b/superset/examples/flights.py @@ -23,7 +23,7 @@ from .helpers import get_example_data, TBL -def load_flights(only_metadata=False, force=False): +def load_flights(only_metadata: bool = False, force: bool = False) -> None: """Loading random time series data from a zip file in the repo""" tbl_name = "flights" database = utils.get_example_database() diff --git a/superset/examples/helpers.py b/superset/examples/helpers.py index 5fac6a0d61e03..9c300bbeba379 100644 --- a/superset/examples/helpers.py +++ b/superset/examples/helpers.py @@ -19,7 +19,7 @@ import os import zlib from io import BytesIO -from typing import Set +from typing import Any, Dict, List, Set from urllib import request from superset import app, db @@ -41,7 +41,7 @@ misc_dash_slices: Set[str] = set() # slices assembled in a 'Misc Chart' dashboard -def update_slice_ids(layout_dict, slices): +def update_slice_ids(layout_dict: Dict[Any, Any], slices: List[Slice]) -> None: charts = [ component for component in layout_dict.values() @@ -53,7 +53,7 @@ def update_slice_ids(layout_dict, slices): chart_component["meta"]["chartId"] = int(slices[i].id) -def merge_slice(slc): +def merge_slice(slc: Slice) -> None: o = db.session.query(Slice).filter_by(slice_name=slc.slice_name).first() if o: db.session.delete(o) @@ -61,13 +61,15 @@ def merge_slice(slc): db.session.commit() -def get_slice_json(defaults, **kwargs): +def get_slice_json(defaults: Dict[Any, Any], **kwargs: Any) -> str: d = defaults.copy() d.update(kwargs) return json.dumps(d, indent=4, sort_keys=True) -def get_example_data(filepath, is_gzip=True, make_bytes=False): +def get_example_data( + filepath: str, is_gzip: bool = True, make_bytes: bool = False +) -> BytesIO: content = request.urlopen(f"{BASE_URL}{filepath}?raw=true").read() if is_gzip: content = zlib.decompress(content, zlib.MAX_WBITS | 16) diff --git a/superset/examples/long_lat.py b/superset/examples/long_lat.py index d90b8b485a7a0..55fce5b1eb22d 100644 --- a/superset/examples/long_lat.py +++ b/superset/examples/long_lat.py @@ -34,7 +34,7 @@ ) -def load_long_lat_data(only_metadata=False, force=False): +def load_long_lat_data(only_metadata: bool = False, force: bool = False) -> None: """Loading lat/long data from a csv file in the repo""" tbl_name = "long_lat" database = utils.get_example_database() diff --git a/superset/examples/misc_dashboard.py b/superset/examples/misc_dashboard.py index 2c90dbbf3a822..a8f282ee431cf 100644 --- a/superset/examples/misc_dashboard.py +++ b/superset/examples/misc_dashboard.py @@ -26,7 +26,7 @@ DASH_SLUG = "misc_charts" -def load_misc_dashboard(): +def load_misc_dashboard() -> None: """Loading a dashboard featuring misc charts""" print("Creating the dashboard") diff --git a/superset/examples/multi_line.py b/superset/examples/multi_line.py index b04db8a620892..1887fd09069e7 100644 --- a/superset/examples/multi_line.py +++ b/superset/examples/multi_line.py @@ -24,7 +24,7 @@ from .world_bank import load_world_bank_health_n_pop -def load_multi_line(only_metadata=False): +def load_multi_line(only_metadata: bool = False) -> None: load_world_bank_health_n_pop(only_metadata) load_birth_names(only_metadata) ids = [ diff --git a/superset/examples/multiformat_time_series.py b/superset/examples/multiformat_time_series.py index 97a7d9566cf7e..dfc36eed65c2e 100644 --- a/superset/examples/multiformat_time_series.py +++ b/superset/examples/multiformat_time_series.py @@ -14,6 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +from typing import Dict, Optional, Tuple import pandas as pd from sqlalchemy import BigInteger, Date, DateTime, String @@ -32,7 +33,9 @@ ) -def load_multiformat_time_series(only_metadata=False, force=False): +def load_multiformat_time_series( + only_metadata: bool = False, force: bool = False +) -> None: """Loading time series data from a zip file in the repo""" tbl_name = "multiformat_time_series" database = get_example_database() @@ -70,15 +73,15 @@ def load_multiformat_time_series(only_metadata=False, force=False): obj = TBL(table_name=tbl_name) obj.main_dttm_col = "ds" obj.database = database - dttm_and_expr_dict = { - "ds": [None, None], - "ds2": [None, None], - "epoch_s": ["epoch_s", None], - "epoch_ms": ["epoch_ms", None], - "string2": ["%Y%m%d-%H%M%S", None], - "string1": ["%Y-%m-%d^%H:%M:%S", None], - "string0": ["%Y-%m-%d %H:%M:%S.%f", None], - "string3": ["%Y/%m/%d%H:%M:%S.%f", None], + dttm_and_expr_dict: Dict[str, Tuple[Optional[str], None]] = { + "ds": (None, None), + "ds2": (None, None), + "epoch_s": ("epoch_s", None), + "epoch_ms": ("epoch_ms", None), + "string2": ("%Y%m%d-%H%M%S", None), + "string1": ("%Y-%m-%d^%H:%M:%S", None), + "string0": ("%Y-%m-%d %H:%M:%S.%f", None), + "string3": ("%Y/%m/%d%H:%M:%S.%f", None), } for col in obj.columns: dttm_and_expr = dttm_and_expr_dict[col.column_name] diff --git a/superset/examples/paris.py b/superset/examples/paris.py index 482ada6115e68..a9ea87184615f 100644 --- a/superset/examples/paris.py +++ b/superset/examples/paris.py @@ -25,7 +25,7 @@ from .helpers import get_example_data, TBL -def load_paris_iris_geojson(only_metadata=False, force=False): +def load_paris_iris_geojson(only_metadata: bool = False, force: bool = False) -> None: tbl_name = "paris_iris_mapping" database = utils.get_example_database() table_exists = database.has_table_by_name(tbl_name) diff --git a/superset/examples/random_time_series.py b/superset/examples/random_time_series.py index 7ce865cbd33fd..41eb98e2c0c52 100644 --- a/superset/examples/random_time_series.py +++ b/superset/examples/random_time_series.py @@ -25,7 +25,9 @@ from .helpers import config, get_example_data, get_slice_json, merge_slice, TBL -def load_random_time_series_data(only_metadata=False, force=False): +def load_random_time_series_data( + only_metadata: bool = False, force: bool = False +) -> None: """Loading random time series data from a zip file in the repo""" tbl_name = "random_time_series" database = utils.get_example_database() diff --git a/superset/examples/sf_population_polygons.py b/superset/examples/sf_population_polygons.py index a4281d9225bca..952bdccc5399d 100644 --- a/superset/examples/sf_population_polygons.py +++ b/superset/examples/sf_population_polygons.py @@ -25,7 +25,9 @@ from .helpers import get_example_data, TBL -def load_sf_population_polygons(only_metadata=False, force=False): +def load_sf_population_polygons( + only_metadata: bool = False, force: bool = False +) -> None: tbl_name = "sf_population_polygons" database = utils.get_example_database() table_exists = database.has_table_by_name(tbl_name) diff --git a/superset/examples/tabbed_dashboard.py b/superset/examples/tabbed_dashboard.py index a35087767fe2a..1eaed48f25e70 100644 --- a/superset/examples/tabbed_dashboard.py +++ b/superset/examples/tabbed_dashboard.py @@ -25,7 +25,7 @@ from .helpers import update_slice_ids -def load_tabbed_dashboard(_=False): +def load_tabbed_dashboard(_: bool = False) -> None: """Creating a tabbed dashboard""" print("Creating a dashboard with nested tabs") diff --git a/superset/examples/unicode_test_data.py b/superset/examples/unicode_test_data.py index d48dc34dac650..70d5e72d3ae91 100644 --- a/superset/examples/unicode_test_data.py +++ b/superset/examples/unicode_test_data.py @@ -36,7 +36,7 @@ ) -def load_unicode_test_data(only_metadata=False, force=False): +def load_unicode_test_data(only_metadata: bool = False, force: bool = False) -> None: """Loading unicode test dataset from a csv file in the repo""" tbl_name = "unicode_test" database = utils.get_example_database() diff --git a/superset/examples/world_bank.py b/superset/examples/world_bank.py index b30d07fe0f320..0136ed9546785 100644 --- a/superset/examples/world_bank.py +++ b/superset/examples/world_bank.py @@ -41,9 +41,9 @@ ) -def load_world_bank_health_n_pop( - only_metadata=False, force=False -): # pylint: disable=too-many-locals +def load_world_bank_health_n_pop( # pylint: disable=too-many-locals + only_metadata: bool = False, force: bool = False +) -> None: """Loads the world bank health dataset, slices and a dashboard""" tbl_name = "wb_health_population" database = utils.get_example_database() diff --git a/superset/utils/core.py b/superset/utils/core.py index 778c71eac713d..e1ab257fcfa4a 100644 --- a/superset/utils/core.py +++ b/superset/utils/core.py @@ -37,7 +37,18 @@ from email.utils import formatdate from enum import Enum from time import struct_time -from typing import Any, Dict, Iterator, List, NamedTuple, Optional, Set, Tuple, Union +from typing import ( + Any, + Dict, + Iterator, + List, + NamedTuple, + Optional, + Set, + Tuple, + TYPE_CHECKING, + Union, +) from urllib.parse import unquote_plus import bleach @@ -72,6 +83,9 @@ except ImportError: pass +if TYPE_CHECKING: + from superset.models.core import Database + logging.getLogger("MARKDOWN").setLevel(logging.INFO) logger = logging.getLogger(__name__) @@ -944,7 +958,7 @@ def get_or_create_db(database_name, sqlalchemy_uri, *args, **kwargs): return database -def get_example_database(): +def get_example_database() -> "Database": from superset import conf db_uri = conf.get("SQLALCHEMY_EXAMPLES_URI") or conf.get("SQLALCHEMY_DATABASE_URI") @@ -1057,11 +1071,15 @@ def get_since_until( else: rel, num, grain = time_range.split() if rel == "Last": - since = relative_start - relativedelta(**{grain: int(num)}) # type: ignore + since = relative_start - relativedelta( # type: ignore + **{grain: int(num)} # type: ignore + ) until = relative_end else: # rel == 'Next' since = relative_start - until = relative_end + relativedelta(**{grain: int(num)}) # type: ignore + until = relative_end + relativedelta( # type: ignore + **{grain: int(num)} # type: ignore + ) else: since = since or "" if since: diff --git a/superset/viz.py b/superset/viz.py index f4890ec69b043..5a9a6ac2bde02 100644 --- a/superset/viz.py +++ b/superset/viz.py @@ -1875,8 +1875,8 @@ def get_data(self, df: pd.DataFrame) -> VizData: for row in d: country = None if isinstance(row["country"], str): - country = countries.get(fd.get("country_fieldtype"), row["country"]) - + if "country_fieldtype" in fd: + country = countries.get(fd["country_fieldtype"], row["country"]) if country: row["country"] = country["cca3"] row["latitude"] = country["lat"]