chore: Embrace the walrus operator

.pre-commit-config.yaml

@@ -19,6 +19,10 @@ repos:
  rev: 5.12.0
  hooks:
  - id: isort
+ - repo: https://github.com/MarcoGorelli/auto-walrus
+ rev: v0.2.2
+ hooks:
+ - id: auto-walrus
  - repo: https://github.com/pre-commit/mirrors-mypy
  rev: v1.3.0
  hooks:

superset/charts/api.py

@@ -649,8 +649,7 @@ def screenshot(self, pk: int, digest: str) -> WerkzeugResponse:
  return self.response_404()
 
  # fetch the chart screenshot using the current user and cache if set
- img = ChartScreenshot.get_from_cache_key(thumbnail_cache, digest)
- if img:
+ if img := ChartScreenshot.get_from_cache_key(thumbnail_cache, digest):
  return Response(
  FileWrapper(img), mimetype="image/png", direct_passthrough=True
  )
@@ -783,7 +782,6 @@ def export(self, **kwargs: Any) -> Response:
  500:
  $ref: '#/components/responses/500'
  """
- token = request.args.get("token")
  requested_ids = kwargs["rison"]
  timestamp = datetime.now().strftime("%Y%m%dT%H%M%S")
  root = f"chart_export_{timestamp}"
@@ -805,7 +803,7 @@ def export(self, **kwargs: Any) -> Response:
  as_attachment=True,
  download_name=filename,
  )
- if token:
+ if token := request.args.get("token"):
  response.set_cookie(token, "done", max_age=600)
  return response
 

superset/charts/commands/bulk_delete.py

@@ -55,8 +55,7 @@ def validate(self) -> None:
  if not self._models or len(self._models) != len(self._model_ids):
  raise ChartNotFoundError()
  # Check there are no associated ReportSchedules
- reports = ReportScheduleDAO.find_by_chart_ids(self._model_ids)
- if reports:
+ if reports := ReportScheduleDAO.find_by_chart_ids(self._model_ids):
  report_names = [report.name for report in reports]
  raise ChartBulkDeleteFailedReportsExistError(
  _("There are associated alerts or reports: %s" % ",".join(report_names))

superset/charts/commands/delete.py

@@ -64,8 +64,7 @@ def validate(self) -> None:
  if not self._model:
  raise ChartNotFoundError()
  # Check there are no associated ReportSchedules
- reports = ReportScheduleDAO.find_by_chart_id(self._model_id)
- if reports:
+ if reports := ReportScheduleDAO.find_by_chart_id(self._model_id):
  report_names = [report.name for report in reports]
  raise ChartDeleteFailedReportsExistError(
  _("There are associated alerts or reports: %s" % ",".join(report_names))

superset/common/query_actions.py

@@ -221,8 +221,7 @@ def get_query_results(
  :raises QueryObjectValidationError: if an unsupported result type is requested
  :return: JSON serializable result payload
  """
- result_func = _result_type_functions.get(result_type)
- if result_func:
+ if result_func := _result_type_functions.get(result_type):
  return result_func(query_context, query_obj, force_cached)
  raise QueryObjectValidationError(
  _("Invalid result type: %(result_type)s", result_type=result_type)

superset/common/query_context_factory.py

@@ -125,10 +125,9 @@ def _apply_granularity(
  for column in datasource.columns
  if (column["is_dttm"] if isinstance(column, dict) else column.is_dttm)
  }
- granularity = query_object.granularity
  x_axis = form_data and form_data.get("x_axis")
 
- if granularity:
+ if granularity := query_object.granularity:
  filter_to_remove = None
  if x_axis and x_axis in temporal_columns:
  filter_to_remove = x_axis

superset/common/query_context_processor.py

@@ -500,8 +500,7 @@ def get_payload(
  return return_value
 
  def get_cache_timeout(self) -> int:
- cache_timeout_rv = self._query_context.get_cache_timeout()
- if cache_timeout_rv:
+ if cache_timeout_rv := self._query_context.get_cache_timeout():
  return cache_timeout_rv
  if (
  data_cache_timeout := config["DATA_CACHE_CONFIG"].get(

superset/common/utils/query_cache_manager.py

@@ -148,8 +148,7 @@ def get(
  if not key or not _cache[region] or force_query:
  return query_cache
 
- cache_value = _cache[region].get(key)
- if cache_value:
+ if cache_value := _cache[region].get(key):
  logger.debug("Cache key: %s", key)
  stats_logger.incr("loading_from_cache")
  try:

superset/connectors/sqla/models.py

@@ -993,11 +993,10 @@ def adhoc_column_to_sqla( # pylint: disable=too-many-locals
  schema=self.schema,
  template_processor=template_processor,
  )
- col_in_metadata = self.get_column(expression)
  time_grain = col.get("timeGrain")
  has_timegrain = col.get("columnType") == "BASE_AXIS" and time_grain
  is_dttm = False
- if col_in_metadata:
+ if col_in_metadata := self.get_column(expression):
  sqla_column = col_in_metadata.get_sqla_col(
  template_processor=template_processor
  )

superset/dashboards/commands/bulk_delete.py

@@ -56,8 +56,7 @@ def validate(self) -> None:
  if not self._models or len(self._models) != len(self._model_ids):
  raise DashboardNotFoundError()
  # Check there are no associated ReportSchedules
- reports = ReportScheduleDAO.find_by_dashboard_ids(self._model_ids)
- if reports:
+ if reports := ReportScheduleDAO.find_by_dashboard_ids(self._model_ids):
  report_names = [report.name for report in reports]
  raise DashboardBulkDeleteFailedReportsExistError(
  _("There are associated alerts or reports: %s" % ",".join(report_names))

superset/dashboards/commands/delete.py

@@ -57,8 +57,7 @@ def validate(self) -> None:
  if not self._model:
  raise DashboardNotFoundError()
  # Check there are no associated ReportSchedules
- reports = ReportScheduleDAO.find_by_dashboard_id(self._model_id)
- if reports:
+ if reports := ReportScheduleDAO.find_by_dashboard_id(self._model_id):
  report_names = [report.name for report in reports]
  raise DashboardDeleteFailedReportsExistError(
  _("There are associated alerts or reports: %s" % ",".join(report_names))

superset/dashboards/dao.py

@@ -200,11 +200,10 @@ def set_dash_metadata( # pylint: disable=too-many-locals
  old_to_new_slice_ids: Optional[Dict[int, int]] = None,
  commit: bool = False,
  ) -> Dashboard:
- positions = data.get("positions")
  new_filter_scopes = {}
  md = dashboard.params_dict
 
- if positions is not None:
+ if (positions := data.get("positions")) is not None:
  # find slices in the position data
  slice_ids = [
  value.get("meta", {}).get("chartId")

superset/databases/api.py

@@ -1036,7 +1036,6 @@ def export(self, **kwargs: Any) -> Response:
  500:
  $ref: '#/components/responses/500'
  """
- token = request.args.get("token")
  requested_ids = kwargs["rison"]
  timestamp = datetime.now().strftime("%Y%m%dT%H%M%S")
  root = f"database_export_{timestamp}"
@@ -1060,7 +1059,7 @@ def export(self, **kwargs: Any) -> Response:
  as_attachment=True,
  download_name=filename,
  )
- if token:
+ if token := request.args.get("token"):
  response.set_cookie(token, "done", max_age=600)
  return response
 

superset/databases/commands/delete.py

@@ -55,9 +55,8 @@ def validate(self) -> None:
  if not self._model:
  raise DatabaseNotFoundError()
  # Check there are no associated ReportSchedules
- reports = ReportScheduleDAO.find_by_database_id(self._model_id)
 
- if reports:
+ if reports := ReportScheduleDAO.find_by_database_id(self._model_id):
  report_names = [report.name for report in reports]
  raise DatabaseDeleteFailedReportsExistError(
  _("There are associated alerts or reports: %s" % ",".join(report_names))

superset/databases/commands/test_connection.py

@@ -228,6 +228,5 @@ def ping(engine: Engine) -> bool:
  raise DatabaseTestConnectionUnexpectedError(errors) from ex
 
  def validate(self) -> None:
- database_name = self._properties.get("database_name")
- if database_name is not None:
+ if (database_name := self._properties.get("database_name")) is not None:
  self._model = DatabaseDAO.get_database_by_name(database_name)

superset/databases/commands/validate.py

@@ -128,6 +128,5 @@ def run(self) -> None:
  )
 
  def validate(self) -> None:
- database_id = self._properties.get("id")
- if database_id is not None:
+ if (database_id := self._properties.get("id")) is not None:
  self._model = DatabaseDAO.find_by_id(database_id)

superset/datasets/api.py

@@ -977,8 +977,7 @@ def get_or_create_dataset(self) -> Response:
  return self.response(400, message=ex.messages)
  table_name = body["table_name"]
  database_id = body["database_id"]
- table = DatasetDAO.get_table_by_name(database_id, table_name)
- if table:
+ if table := DatasetDAO.get_table_by_name(database_id, table_name):
  return self.response(200, result={"table_id": table.id})
 
  body["database"] = database_id

superset/datasets/commands/importers/v1/utils.py

@@ -62,8 +62,7 @@ def get_sqla_type(native_type: str) -> VisitableType:
  if native_type.upper() in type_map:
  return type_map[native_type.upper()]
 
- match = VARCHAR.match(native_type)
- if match:
+ if match := VARCHAR.match(native_type):
  size = int(match.group(1))
  return String(size)
 

superset/datasets/commands/update.py

@@ -114,13 +114,11 @@ def validate(self) -> None:
  exceptions.append(DatasetEndpointUnsafeValidationError())
 
  # Validate columns
- columns = self._properties.get("columns")
- if columns:
+ if columns := self._properties.get("columns"):
  self._validate_columns(columns, exceptions)
 
  # Validate metrics
- metrics = self._properties.get("metrics")
- if metrics:
+ if metrics := self._properties.get("metrics"):
  self._validate_metrics(metrics, exceptions)
 
  if exceptions:

superset/db_engine_specs/base.py

@@ -1704,8 +1704,7 @@ def get_column_spec( # pylint: disable=unused-argument
  :param source: Type coming from the database table or cursor description
  :return: ColumnSpec object
  """
- col_types = cls.get_column_types(native_type)
- if col_types:
+ if col_types := cls.get_column_types(native_type):
  column_type, generic_type = col_types
  is_dttm = generic_type == GenericDataType.TEMPORAL
  return ColumnSpec(
@@ -1996,9 +1995,8 @@ def validate_parameters(
  required = {"host", "port", "username", "database"}
  parameters = properties.get("parameters", {})
  present = {key for key in parameters if parameters.get(key, ())}
- missing = sorted(required - present)
 
- if missing:
+ if missing := sorted(required - present):
  errors.append(
  SupersetError(
  message=f'One or more parameters are missing: {", ".join(missing)}',

superset/db_engine_specs/bigquery.py

@@ -384,9 +384,8 @@ def df_to_sql(
  }
 
  # Add credentials if they are set on the SQLAlchemy dialect.
- creds = engine.dialect.credentials_info
 
- if creds:
+ if creds := engine.dialect.credentials_info:
  to_gbq_kwargs[
  "credentials"
  ] = service_account.Credentials.from_service_account_info(creds)

superset/db_engine_specs/databricks.py

@@ -285,9 +285,8 @@ def validate_parameters( # type: ignore
  parameters["http_path"] = connect_args.get("http_path")
 
  present = {key for key in parameters if parameters.get(key, ())}
- missing = sorted(required - present)
 
- if missing:
+ if missing := sorted(required - present):
  errors.append(
  SupersetError(
  message=f'One or more parameters are missing: {", ".join(missing)}',

superset/db_engine_specs/presto.py

@@ -1213,8 +1213,7 @@ def extra_table_metadata(
  ) -> Dict[str, Any]:
  metadata = {}
 
- indexes = database.get_indexes(table_name, schema_name)
- if indexes:
+ if indexes := database.get_indexes(table_name, schema_name):
  col_names, latest_parts = cls.latest_partition(
  table_name, schema_name, database, show_first=True
  )
@@ -1278,8 +1277,7 @@ def get_tracking_url(cls, cursor: "Cursor") -> Optional[str]:
  @classmethod
  def handle_cursor(cls, cursor: "Cursor", query: Query, session: Session) -> None:
  """Updates progress information"""
- tracking_url = cls.get_tracking_url(cursor)
- if tracking_url:
+ if tracking_url := cls.get_tracking_url(cursor):
  query.tracking_url = tracking_url
  session.commit()
 

superset/db_engine_specs/snowflake.py

@@ -312,9 +312,8 @@ def validate_parameters(
  }
  parameters = properties.get("parameters", {})
  present = {key for key in parameters if parameters.get(key, ())}
- missing = sorted(required - present)
 
- if missing:
+ if missing := sorted(required - present):
  errors.append(
  SupersetError(
  message=f'One or more parameters are missing: {", ".join(missing)}',

superset/db_engine_specs/trino.py

@@ -57,8 +57,7 @@ def extra_table_metadata(
  ) -> Dict[str, Any]:
  metadata = {}
 
- indexes = database.get_indexes(table_name, schema_name)
- if indexes:
+ if indexes := database.get_indexes(table_name, schema_name):
  col_names, latest_parts = cls.latest_partition(
  table_name, schema_name, database, show_first=True
  )
@@ -150,8 +149,7 @@ def get_tracking_url(cls, cursor: Cursor) -> Optional[str]:
 
  @classmethod
  def handle_cursor(cls, cursor: Cursor, query: Query, session: Session) -> None:
- tracking_url = cls.get_tracking_url(cursor)
- if tracking_url:
+ if tracking_url := cls.get_tracking_url(cursor):
  query.tracking_url = tracking_url
 
  # Adds the executed query id to the extra payload so the query can be cancelled

superset/errors.py

@@ -211,8 +211,7 @@ def __post_init__(self) -> None:
  Mutates the extra params with user facing error codes that map to backend
  errors.
  """
- issue_codes = ERROR_TYPES_TO_ISSUE_CODES_MAPPING.get(self.error_type)
- if issue_codes:
+ if issue_codes := ERROR_TYPES_TO_ISSUE_CODES_MAPPING.get(self.error_type):
  self.extra = self.extra or {}
  self.extra.update(
  {

superset/initialization/__init__.py

@@ -453,8 +453,7 @@ def init_app_in_ctx(self) -> None:
 
  # Hook that provides administrators a handle on the Flask APP
  # after initialization
- flask_app_mutator = self.config["FLASK_APP_MUTATOR"]
- if flask_app_mutator:
+ if flask_app_mutator := self.config["FLASK_APP_MUTATOR"]:
  flask_app_mutator(self.superset_app)
 
  if feature_flag_manager.is_feature_enabled("TAGGING_SYSTEM"):

superset/migrations/env.py

@@ -103,16 +103,15 @@ def process_revision_directives( # pylint: disable=redefined-outer-name, unused
  kwargs = {}
  if engine.name in ("sqlite", "mysql"):
  kwargs = {"transaction_per_migration": True, "transactional_ddl": True}
- configure_args = current_app.extensions["migrate"].configure_args
- if configure_args:
+ if configure_args := current_app.extensions["migrate"].configure_args:
  kwargs.update(configure_args)
 
  context.configure(
  connection=connection,
  target_metadata=target_metadata,
  # compare_type=True,
  process_revision_directives=process_revision_directives,
- **kwargs
+ **kwargs,
  )
 
  try:

superset/migrations/shared/migrate_viz/base.py

@@ -130,17 +130,15 @@ def upgrade_slice(cls, slc: Slice) -> Slice:
  # only backup params
  slc.params = json.dumps({**clz.data, FORM_DATA_BAK_FIELD_NAME: form_data_bak})
 
- query_context = try_load_json(slc.query_context)
- if "form_data" in query_context:
+ if "form_data" in (query_context := try_load_json(slc.query_context)):
  query_context["form_data"] = clz.data
  slc.query_context = json.dumps(query_context)
  return slc
 
  @classmethod
  def downgrade_slice(cls, slc: Slice) -> Slice:
  form_data = try_load_json(slc.params)
- form_data_bak = form_data.get(FORM_DATA_BAK_FIELD_NAME, {})
- if "viz_type" in form_data_bak:
+ if "viz_type" in (form_data_bak := form_data.get(FORM_DATA_BAK_FIELD_NAME, {})):
  slc.params = json.dumps(form_data_bak)
  slc.viz_type = form_data_bak.get("viz_type")
  query_context = try_load_json(slc.query_context)

superset/migrations/shared/migrate_viz/processors.py

@@ -40,16 +40,14 @@ def _pre_action(self) -> None:
  if self.data.get("contribution"):
  self.data["contributionMode"] = "row"
 
- stacked = self.data.get("stacked_style")
- if stacked:
+ if stacked := self.data.get("stacked_style"):
  stacked_map = {
  "expand": "Expand",
  "stack": "Stack",
  }
  self.data["show_extra_controls"] = True
  self.data["stack"] = stacked_map.get(stacked)
 
- x_axis_label = self.data.get("x_axis_label")
- if x_axis_label:
+ if x_axis_label := self.data.get("x_axis_label"):
  self.data["x_axis_title"] = x_axis_label
  self.data["x_axis_title_margin"] = 30