Skip to content

Commit f33e29e

Browse files
committed
Add mypy plugin for decorators.
* Preserve types for decorators that don't change signatures * Augment types for decorators that change signatures * Drop redundant type checks
1 parent 77412ac commit f33e29e

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

53 files changed

+624
-1486
lines changed

UPDATING.md

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,35 @@ https://developers.google.com/style/inclusive-documentation
6262
6363
-->
6464

65+
### Added mypy plugin to preserve types of decorated functions
66+
67+
Mypy currently doesn't support precise type information for decorated
68+
functions; see https://github.com/python/mypy/issues/3157 for details.
69+
To preserve precise type definitions for decorated functions, we now
70+
include a mypy plugin to preserve precise type definitions for decorated
71+
functions. To use the plugin, update your setup.cfg:
72+
73+
```
74+
[mypy]
75+
plugins =
76+
airflow.mypy.plugin.decorators
77+
```
78+
79+
### Use project_id argument consistently across GCP hooks and operators
80+
81+
- Changed order of arguments in DataflowHook.start_python_dataflow. Uses
82+
with positional arguments may break.
83+
- Changed order of arguments in DataflowHook.is_job_dataflow_running. Uses
84+
with positional arguments may break.
85+
- Changed order of arguments in DataflowHook.cancel_job. Uses
86+
with positional arguments may break.
87+
- Added optional project_id argument to DataflowCreateJavaJobOperator
88+
constructor.
89+
- Added optional project_id argument to DataflowTemplatedJobStartOperator
90+
constructor.
91+
- Added optional project_id argument to DataflowCreatePythonJobOperator
92+
constructor.
93+
6594
### Rename pool statsd metrics
6695

6796
Used slot has been renamed to running slot to make the name self-explanatory

airflow/mypy/__init__.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
#
2+
# Licensed to the Apache Software Foundation (ASF) under one
3+
# or more contributor license agreements. See the NOTICE file
4+
# distributed with this work for additional information
5+
# regarding copyright ownership. The ASF licenses this file
6+
# to you under the Apache License, Version 2.0 (the
7+
# "License"); you may not use this file except in compliance
8+
# with the License. You may obtain a copy of the License at
9+
#
10+
# http://www.apache.org/licenses/LICENSE-2.0
11+
#
12+
# Unless required by applicable law or agreed to in writing,
13+
# software distributed under the License is distributed on an
14+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
# KIND, either express or implied. See the License for the
16+
# specific language governing permissions and limitations
17+
# under the License.

airflow/mypy/plugin/__init__.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
#
2+
# Licensed to the Apache Software Foundation (ASF) under one
3+
# or more contributor license agreements. See the NOTICE file
4+
# distributed with this work for additional information
5+
# regarding copyright ownership. The ASF licenses this file
6+
# to you under the Apache License, Version 2.0 (the
7+
# "License"); you may not use this file except in compliance
8+
# with the License. You may obtain a copy of the License at
9+
#
10+
# http://www.apache.org/licenses/LICENSE-2.0
11+
#
12+
# Unless required by applicable law or agreed to in writing,
13+
# software distributed under the License is distributed on an
14+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
# KIND, either express or implied. See the License for the
16+
# specific language governing permissions and limitations
17+
# under the License.

airflow/mypy/plugin/decorators.py

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
#
2+
# Licensed to the Apache Software Foundation (ASF) under one
3+
# or more contributor license agreements. See the NOTICE file
4+
# distributed with this work for additional information
5+
# regarding copyright ownership. The ASF licenses this file
6+
# to you under the Apache License, Version 2.0 (the
7+
# "License"); you may not use this file except in compliance
8+
# with the License. You may obtain a copy of the License at
9+
#
10+
# http://www.apache.org/licenses/LICENSE-2.0
11+
#
12+
# Unless required by applicable law or agreed to in writing,
13+
# software distributed under the License is distributed on an
14+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
# KIND, either express or implied. See the License for the
16+
# specific language governing permissions and limitations
17+
# under the License.
18+
19+
import copy
20+
import functools
21+
from typing import List
22+
23+
from mypy.nodes import ARG_NAMED_OPT # pylint: disable=no-name-in-module
24+
from mypy.plugin import FunctionContext, Plugin # pylint: disable=no-name-in-module
25+
from mypy.types import CallableType, NoneType, UnionType # pylint: disable=no-name-in-module
26+
27+
TYPED_DECORATORS = {
28+
"fallback_to_default_project_id of GoogleBaseHook": ["project_id"],
29+
"airflow.providers.google.cloud.hooks.dataflow._fallback_to_project_id_from_variables": ["project_id"],
30+
"provide_gcp_credential_file of GoogleBaseHook": [],
31+
}
32+
33+
34+
class TypedDecoratorPlugin(Plugin):
35+
"""Mypy plugin for typed decorators."""
36+
def get_function_hook(self, fullname: str):
37+
"""Check for known typed decorators by name."""
38+
if fullname in TYPED_DECORATORS:
39+
return functools.partial(
40+
_analyze_decorator,
41+
provided_arguments=TYPED_DECORATORS[fullname],
42+
)
43+
return None
44+
45+
46+
def _analyze_decorator(function_ctx: FunctionContext, provided_arguments: List[str]):
47+
if not isinstance(function_ctx.arg_types[0][0], CallableType):
48+
return function_ctx.default_return_type
49+
if not isinstance(function_ctx.default_return_type, CallableType):
50+
return function_ctx.default_return_type
51+
return _change_decorator_function_type(
52+
function_ctx.arg_types[0][0],
53+
function_ctx.default_return_type,
54+
provided_arguments,
55+
)
56+
57+
58+
def _change_decorator_function_type(
59+
decorated: CallableType,
60+
decorator: CallableType,
61+
provided_arguments: List[str],
62+
) -> CallableType:
63+
decorator.arg_kinds = decorated.arg_kinds
64+
decorator.arg_names = decorated.arg_names
65+
66+
# Mark provided arguments as optional
67+
decorator.arg_types = copy.copy(decorated.arg_types)
68+
for argument in provided_arguments:
69+
index = decorated.arg_names.index(argument)
70+
decorated_type = decorated.arg_types[index]
71+
decorator.arg_types[index] = UnionType.make_union([decorated_type, NoneType()])
72+
decorated.arg_kinds[index] = ARG_NAMED_OPT
73+
74+
return decorator
75+
76+
77+
def plugin(version: str): # pylint: disable=unused-argument
78+
"""Mypy plugin entrypoint."""
79+
return TypedDecoratorPlugin

airflow/providers/google/cloud/example_dags/example_automl_tables.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -139,7 +139,6 @@ def get_target_column_spec(columns_specs: List[Dict], column_name: str) -> str:
139139
task_id="update_dataset_task",
140140
dataset=update,
141141
location=GCP_AUTOML_LOCATION,
142-
project_id=GCP_PROJECT_ID,
143142
)
144143
# [END howto_operator_automl_update_dataset]
145144

0 commit comments

Comments
 (0)