Enable Black - Python Auto Formmatter (#9550)

apache · Nov 3, 2020 · 4e8f9cc · 4e8f9cc
1 parent 1dc7099
commit 4e8f9cc
Show file tree

Hide file tree

Showing 1,070 changed files with 15,413 additions and 15,140 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -159,8 +159,7 @@ repos:
  rev: 20.8b1
  hooks:
  - id: black
- files: api_connexion/.*\.py|.*providers.*\.py|^chart/tests/.*\.py
- exclude: .*kubernetes_pod\.py|.*google/common/hooks/base_google\.py$
+ exclude: .*kubernetes_pod\.py|.*google/common/hooks/base_google\.py$|^airflow/configuration.py$
  args: [--config=./pyproject.toml]
  - repo: https://github.com/pre-commit/pre-commit-hooks
  rev: v3.3.0
@@ -203,7 +202,7 @@ repos:
  name: Run isort to sort imports
  types: [python]
  # To keep consistent with the global isort skip config defined in setup.cfg
- exclude: ^build/.*$|^.tox/.*$|^venv/.*$|.*api_connexion/.*\.py|.*providers.*\.py
+ exclude: ^build/.*$|^.tox/.*$|^venv/.*$
  - repo: https://github.com/pycqa/pydocstyle
  rev: 5.1.1
  hooks:

diff --git a/airflow/__init__.py b/airflow/__init__.py
@@ -55,15 +55,18 @@ def __getattr__(name):
  # PEP-562: Lazy loaded attributes on python modules
  if name == "DAG":
  from airflow.models.dag import DAG # pylint: disable=redefined-outer-name
+
  return DAG
  if name == "AirflowException":
  from airflow.exceptions import AirflowException # pylint: disable=redefined-outer-name
+
  return AirflowException
  raise AttributeError(f"module {__name__} has no attribute {name}")
 
 
 if not settings.LAZY_LOAD_PLUGINS:
  from airflow import plugins_manager
+
  plugins_manager.ensure_plugins_loaded()
 
 

diff --git a/airflow/api/__init__.py b/airflow/api/__init__.py
@@ -38,8 +38,5 @@ def load_auth():
  log.info("Loaded API auth backend: %s", auth_backend)
  return auth_backend
  except ImportError as err:
- log.critical(
- "Cannot import %s for API authentication due to: %s",
- auth_backend, err
- )
+ log.critical("Cannot import %s for API authentication due to: %s", auth_backend, err)
  raise AirflowException(err)
diff --git a/airflow/api/auth/backend/basic_auth.py b/airflow/api/auth/backend/basic_auth.py
@@ -53,13 +53,12 @@ def auth_current_user() -> Optional[User]:
 
 def requires_authentication(function: T):
  """Decorator for functions that require authentication"""
+
  @wraps(function)
  def decorated(*args, **kwargs):
  if auth_current_user() is not None:
  return function(*args, **kwargs)
  else:
- return Response(
- "Unauthorized", 401, {"WWW-Authenticate": "Basic"}
- )
+ return Response("Unauthorized", 401, {"WWW-Authenticate": "Basic"})
 
  return cast(T, decorated)
diff --git a/airflow/api/auth/backend/default.py b/airflow/api/auth/backend/default.py
@@ -33,6 +33,7 @@ def init_app(_):
 
 def requires_authentication(function: T):
  """Decorator for functions that require authentication"""
+
  @wraps(function)
  def decorated(*args, **kwargs):
  return function(*args, **kwargs)

diff --git a/airflow/api/auth/backend/kerberos_auth.py b/airflow/api/auth/backend/kerberos_auth.py
@@ -132,6 +132,7 @@ def _gssapi_authenticate(token):
 
 def requires_authentication(function: T):
  """Decorator for functions that require authentication with Kerberos"""
+
  @wraps(function)
  def decorated(*args, **kwargs):
  header = request.headers.get("Authorization")
@@ -144,11 +145,11 @@ def decorated(*args, **kwargs):
  response = function(*args, **kwargs)
  response = make_response(response)
  if ctx.kerberos_token is not None:
- response.headers['WWW-Authenticate'] = ' '.join(['negotiate',
- ctx.kerberos_token])
+ response.headers['WWW-Authenticate'] = ' '.join(['negotiate', ctx.kerberos_token])
 
  return response
  if return_code != kerberos.AUTH_GSS_CONTINUE:
  return _forbidden()
  return _unauthorized()
+
  return cast(T, decorated)
diff --git a/airflow/api/client/__init__.py b/airflow/api/client/__init__.py
@@ -35,6 +35,6 @@ def get_current_api_client() -> Client:
  api_client = api_module.Client(
  api_base_url=conf.get('cli', 'endpoint_url'),
  auth=getattr(auth_backend, 'CLIENT_AUTH', None),
- session=session
+ session=session,
  )
  return api_client
diff --git a/airflow/api/client/json_client.py b/airflow/api/client/json_client.py
@@ -45,12 +45,15 @@ def _request(self, url, method='GET', json=None):
  def trigger_dag(self, dag_id, run_id=None, conf=None, execution_date=None):
  endpoint = f'/api/experimental/dags/{dag_id}/dag_runs'
  url = urljoin(self._api_base_url, endpoint)
- data = self._request(url, method='POST',
- json={
- "run_id": run_id,
- "conf": conf,
- "execution_date": execution_date,
- })
+ data = self._request(
+ url,
+ method='POST',
+ json={
+ "run_id": run_id,
+ "conf": conf,
+ "execution_date": execution_date,
+ },
+ )
  return data['message']
 
  def delete_dag(self, dag_id):
@@ -74,12 +77,15 @@ def get_pools(self):
  def create_pool(self, name, slots, description):
  endpoint = '/api/experimental/pools'
  url = urljoin(self._api_base_url, endpoint)
- pool = self._request(url, method='POST',
- json={
- 'name': name,
- 'slots': slots,
- 'description': description,
- })
+ pool = self._request(
+ url,
+ method='POST',
+ json={
+ 'name': name,
+ 'slots': slots,
+ 'description': description,
+ },
+ )
  return pool['pool'], pool['slots'], pool['description']
 
  def delete_pool(self, name):

diff --git a/airflow/api/client/local_client.py b/airflow/api/client/local_client.py
@@ -26,10 +26,9 @@ class Client(api_client.Client):
  """Local API client implementation."""
 
  def trigger_dag(self, dag_id, run_id=None, conf=None, execution_date=None):
- dag_run = trigger_dag.trigger_dag(dag_id=dag_id,
- run_id=run_id,
- conf=conf,
- execution_date=execution_date)
+ dag_run = trigger_dag.trigger_dag(
+ dag_id=dag_id, run_id=run_id, conf=conf, execution_date=execution_date
+ )
  return f"Created {dag_run}"
 
  def delete_dag(self, dag_id):

diff --git a/airflow/api/common/experimental/__init__.py b/airflow/api/common/experimental/__init__.py
@@ -29,10 +29,7 @@ def check_and_get_dag(dag_id: str, task_id: Optional[str] = None) -> DagModel:
  if dag_model is None:
  raise DagNotFound(f"Dag id {dag_id} not found in DagModel")
 
- dagbag = DagBag(
- dag_folder=dag_model.fileloc,
- read_dags_from_db=True
- )
+ dagbag = DagBag(dag_folder=dag_model.fileloc, read_dags_from_db=True)
  dag = dagbag.get_dag(dag_id)
  if not dag:
  error_message = f"Dag id {dag_id} not found"
@@ -47,7 +44,6 @@ def check_and_get_dagrun(dag: DagModel, execution_date: datetime) -> DagRun:
  """Get DagRun object and check that it exists"""
  dagrun = dag.get_dagrun(execution_date=execution_date)
  if not dagrun:
- error_message = ('Dag Run for date {} not found in dag {}'
- .format(execution_date, dag.dag_id))
+ error_message = f'Dag Run for date {execution_date} not found in dag {dag.dag_id}'
  raise DagRunNotFound(error_message)
  return dagrun
diff --git a/airflow/api/common/experimental/delete_dag.py b/airflow/api/common/experimental/delete_dag.py
@@ -60,13 +60,14 @@ def delete_dag(dag_id: str, keep_records_in_log: bool = True, session=None) -> i
  if dag.is_subdag:
  parent_dag_id, task_id = dag_id.rsplit(".", 1)
  for model in TaskFail, models.TaskInstance:
- count += session.query(model).filter(model.dag_id == parent_dag_id,
- model.task_id == task_id).delete()
+ count += (
+ session.query(model).filter(model.dag_id == parent_dag_id, model.task_id == task_id).delete()
+ )
 
  # Delete entries in Import Errors table for a deleted DAG
  # This handles the case when the dag_id is changed in the file
- session.query(models.ImportError).filter(
- models.ImportError.filename == dag.fileloc
- ).delete(synchronize_session='fetch')
+ session.query(models.ImportError).filter(models.ImportError.filename == dag.fileloc).delete(
+ synchronize_session='fetch'
+ )
 
  return count
diff --git a/airflow/api/common/experimental/get_dag_runs.py b/airflow/api/common/experimental/get_dag_runs.py
@@ -38,16 +38,16 @@ def get_dag_runs(dag_id: str, state: Optional[str] = None) -> List[Dict[str, Any
  dag_runs = []
  state = state.lower() if state else None
  for run in DagRun.find(dag_id=dag_id, state=state):
- dag_runs.append({
- 'id': run.id,
- 'run_id': run.run_id,
- 'state': run.state,
- 'dag_id': run.dag_id,
- 'execution_date': run.execution_date.isoformat(),
- 'start_date': ((run.start_date or '') and
-   run.start_date.isoformat()),
- 'dag_run_url': url_for('Airflow.graph', dag_id=run.dag_id,
-  execution_date=run.execution_date)
- })
+ dag_runs.append(
+ {
+  'id': run.id,
+  'run_id': run.run_id,
+  'state': run.state,
+  'dag_id': run.dag_id,
+  'execution_date': run.execution_date.isoformat(),
+ 'start_date': ((run.start_date or '') and run.start_date.isoformat()),
+  'dag_run_url': url_for('Airflow.graph', dag_id=run.dag_id, execution_date=run.execution_date),
+ }
+ )
 
  return dag_runs
diff --git a/airflow/api/common/experimental/get_lineage.py b/airflow/api/common/experimental/get_lineage.py
@@ -31,10 +31,12 @@ def get_lineage(dag_id: str, execution_date: datetime.datetime, session=None) ->
  dag = check_and_get_dag(dag_id)
  check_and_get_dagrun(dag, execution_date)
 
- inlets: List[XCom] = XCom.get_many(dag_ids=dag_id, execution_date=execution_date,
- key=PIPELINE_INLETS, session=session).all()
- outlets: List[XCom] = XCom.get_many(dag_ids=dag_id, execution_date=execution_date,
- key=PIPELINE_OUTLETS, session=session).all()
+ inlets: List[XCom] = XCom.get_many(
+ dag_ids=dag_id, execution_date=execution_date, key=PIPELINE_INLETS, session=session
+ ).all()
+ outlets: List[XCom] = XCom.get_many(
+ dag_ids=dag_id, execution_date=execution_date, key=PIPELINE_OUTLETS, session=session
+ ).all()
 
  lineage: Dict[str, Dict[str, Any]] = {}
  for meta in inlets:

diff --git a/airflow/api/common/experimental/get_task_instance.py b/airflow/api/common/experimental/get_task_instance.py
@@ -31,8 +31,7 @@ def get_task_instance(dag_id: str, task_id: str, execution_date: datetime) -> Ta
  # Get task instance object and check that it exists
  task_instance = dagrun.get_task_instance(task_id)
  if not task_instance:
- error_message = ('Task {} instance for date {} not found'
- .format(task_id, execution_date))
+ error_message = f'Task {task_id} instance for date {execution_date} not found'
  raise TaskInstanceNotFound(error_message)
 
  return task_instance
diff --git a/airflow/api/common/experimental/mark_tasks.py b/airflow/api/common/experimental/mark_tasks.py
@@ -69,7 +69,7 @@ def set_state(
  past: bool = False,
  state: str = State.SUCCESS,
  commit: bool = False,
- session=None
+ session=None,
 ): # pylint: disable=too-many-arguments,too-many-locals
  """
  Set the state of a task instance and if needed its relatives. Can set state
@@ -137,33 +137,24 @@ def set_state(
 # Flake and pylint disagree about correct indents here
 def all_subdag_tasks_query(sub_dag_run_ids, session, state, confirmed_dates): # noqa: E123
  """Get *all* tasks of the sub dags"""
- qry_sub_dag = session.query(TaskInstance). \
- filter(
- TaskInstance.dag_id.in_(sub_dag_run_ids),
- TaskInstance.execution_date.in_(confirmed_dates)
- ). \
- filter(
- or_(
- TaskInstance.state.is_(None),
- TaskInstance.state != state
- )
+ qry_sub_dag = (
+ session.query(TaskInstance)
+ .filter(TaskInstance.dag_id.in_(sub_dag_run_ids), TaskInstance.execution_date.in_(confirmed_dates))
+ .filter(or_(TaskInstance.state.is_(None), TaskInstance.state != state))
  ) # noqa: E123
  return qry_sub_dag
 
 
 def get_all_dag_task_query(dag, session, state, task_ids, confirmed_dates):
  """Get all tasks of the main dag that will be affected by a state change"""
- qry_dag = session.query(TaskInstance). \
- filter(
- TaskInstance.dag_id == dag.dag_id,
- TaskInstance.execution_date.in_(confirmed_dates),
- TaskInstance.task_id.in_(task_ids) # noqa: E123
- ). \
- filter(
- or_(
- TaskInstance.state.is_(None),
- TaskInstance.state != state
+ qry_dag = (
+ session.query(TaskInstance)
+ .filter(
+ TaskInstance.dag_id == dag.dag_id,
+ TaskInstance.execution_date.in_(confirmed_dates),
+ TaskInstance.task_id.in_(task_ids), # noqa: E123
  )
+ .filter(or_(TaskInstance.state.is_(None), TaskInstance.state != state))
  )
  return qry_dag
 
@@ -186,10 +177,12 @@ def get_subdag_runs(dag, session, state, task_ids, commit, confirmed_dates):
  # this works as a kind of integrity check
  # it creates missing dag runs for subdag operators,
  # maybe this should be moved to dagrun.verify_integrity
- dag_runs = _create_dagruns(current_task.subdag,
- execution_dates=confirmed_dates,
- state=State.RUNNING,
- run_type=DagRunType.BACKFILL_JOB)
+ dag_runs = _create_dagruns(
+ current_task.subdag,
+ execution_dates=confirmed_dates,
+ state=State.RUNNING,
+ run_type=DagRunType.BACKFILL_JOB,
+ )
 
  verify_dagruns(dag_runs, commit, state, session, current_task)
 
@@ -279,10 +272,9 @@ def _set_dag_run_state(dag_id, execution_date, state, session=None):
  :param state: target state
  :param session: database session
  """
- dag_run = session.query(DagRun).filter(
- DagRun.dag_id == dag_id,
- DagRun.execution_date == execution_date
- ).one()
+ dag_run = (
+ session.query(DagRun).filter(DagRun.dag_id == dag_id, DagRun.execution_date == execution_date).one()
+ )
  dag_run.state = state
  if state == State.RUNNING:
  dag_run.start_date = timezone.utcnow()
@@ -316,8 +308,9 @@ def set_dag_run_state_to_success(dag, execution_date, commit=False, session=None
  # Mark all task instances of the dag run to success.
  for task in dag.tasks:
  task.dag = dag
- return set_state(tasks=dag.tasks, execution_date=execution_date,
- state=State.SUCCESS, commit=commit, session=session)
+ return set_state(
+ tasks=dag.tasks, execution_date=execution_date, state=State.SUCCESS, commit=commit, session=session
+ )
 
 
 @provide_session
@@ -343,10 +336,15 @@ def set_dag_run_state_to_failed(dag, execution_date, commit=False, session=None)
 
  # Mark only RUNNING task instances.
  task_ids = [task.task_id for task in dag.tasks]
- tis = session.query(TaskInstance).filter(
- TaskInstance.dag_id == dag.dag_id,
- TaskInstance.execution_date == execution_date,
- TaskInstance.task_id.in_(task_ids)).filter(TaskInstance.state == State.RUNNING)
+ tis = (
+ session.query(TaskInstance)
+ .filter(
+ TaskInstance.dag_id == dag.dag_id,
+ TaskInstance.execution_date == execution_date,
+ TaskInstance.task_id.in_(task_ids),
+ )
+ .filter(TaskInstance.state == State.RUNNING)
+ )
  task_ids_of_running_tis = [task_instance.task_id for task_instance in tis]
 
  tasks = []
@@ -356,8 +354,9 @@ def set_dag_run_state_to_failed(dag, execution_date, commit=False, session=None)
  task.dag = dag
  tasks.append(task)
 
- return set_state(tasks=tasks, execution_date=execution_date,
- state=State.FAILED, commit=commit, session=session)
+ return set_state(
+ tasks=tasks, execution_date=execution_date, state=State.FAILED, commit=commit, session=session
+ )
 
 
 @provide_session