def _set_module_environment(self, node: astroid.Module) -> None:
"""Method to set environment of a Module node."""
node.type_environment = Environment()
for name in node.globals:
if not any(isinstance(elt, (astroid.ImportFrom, astroid.Import)) for elt in node.globals[name]):
node.type_environment.globals[name] = self.type_constraints.fresh_tvar(node.globals[name][0])
self._populate_local_env(node)
def _set_module_environment(self, node: astroid.Module) -> None:
"""Method to set environment of a Module node."""
node.type_environment = Environment(
globals_={
name: self.type_constraints.fresh_tvar(node.globals[name][0])
for name in node.globals
})
self._populate_local_env(node)
def _set_module_environment(self, node: astroid.Module) -> None:
"""Method to set environment of a Module node."""
node.type_environment = Environment()
for name in node.globals:
if not any(isinstance(elt, (astroid.ImportFrom, astroid.Import)) for elt in node.globals[name]):
new_tvar = self.type_constraints.fresh_tvar(node.globals[name][0])
if any(isinstance(elt, astroid.ClassDef) for elt in node.globals[name]):
self.type_constraints.unify(new_tvar, Type[ForwardRef(name)], node)
node.type_environment.globals[name] = new_tvar
self._populate_local_env(node)
def visit_module(self, module: astroid.Module) -> None:
self.cfgs[module] = ControlFlowGraph()
self._current_cfg = self.cfgs[module]
self._current_block = self._current_cfg.start
module.cfg_block = self._current_cfg.start
for child in module.body:
child.accept(self)
self._current_cfg.link_or_merge(self._current_block, self._current_cfg.end)
def _parse_functions(self, module: astroid.Module) -> None:
"""Parse the function definitions from typeshed."""
for function_def in module.nodes_of_class(astroid.FunctionDef):
in_class = isinstance(function_def.parent, astroid.ClassDef)
if in_class:
tvars = self.classes[function_def.parent.name]['__pyta_tvars']
else:
tvars = []
f_type = parse_annotations(function_def, tvars)
if in_class:
self.classes[function_def.parent.name][function_def.name].extend(f_type)
self.methods[function_def.name].extend(f_type)
else:
self.functions[function_def.name].extend(f_type)
def _parse_classes(self, module: astroid.Module) -> None:
"""Parse the class definitions from typeshed."""
for class_def in module.nodes_of_class(astroid.ClassDef):
tvars = []
self.classes[class_def.name]['__bases'] = []
for base in class_def.bases:
base_type = _node_to_type(base)
self.classes[class_def.name]['__pyta_tvars'] = \
[tv.__name__ for tv in _collect_tvars(base_type)]
self.classes[class_def.name]['__bases'].append(base_type)
self.classes[class_def.name]['__mro'] = [cls.name for cls in class_def.mro()]
for node in (nodes[0] for nodes in class_def.locals.values()
if isinstance(nodes[0], astroid.AssignName) and
isinstance(nodes[0].parent, astroid.AnnAssign)):
self.classes[class_def.name][node.name] = parse_annotations(node, tvars)
def _set_module_environment(self, node: astroid.Module) -> None:
"""Method to set environment of a Module node."""
node.type_environment = Environment()
for name in node.globals:
if not any(
isinstance(elt, (astroid.ImportFrom, astroid.Import))
for elt in node.globals[name]):
new_tvar = self.type_constraints.fresh_tvar(
node.globals[name][0])
if any(
isinstance(elt, astroid.ClassDef)
for elt in node.globals[name]):
self.type_constraints.unify(new_tvar,
Type[ForwardRef(name)], node)
node.type_environment.globals[name] = new_tvar
self._populate_local_env(node)
def _parse_classes(self, module: astroid.Module) -> None:
"""Parse the class definitions from typeshed."""
for class_def in module.nodes_of_class(astroid.ClassDef):
tvars = []
for base in class_def.bases:
if isinstance(base, astroid.Subscript):
gen = base.value.as_string()
tvars = base.slice.as_string().strip('()').replace(" ", "").split(',')
if gen == 'Generic':
self.classes[class_def.name]['__pyta_tvars'] = tvars
for node in (nodes[0] for nodes in class_def.locals.values()
if isinstance(nodes[0], astroid.AssignName) and
isinstance(nodes[0].parent, astroid.AnnAssign)):
self.classes[class_def.name][node.name] = [
parse_annotations(node, tvars)
]
def visit_module(self, node: astroid.Module) -> None:
node.inf_type = NoType()
def visit_module(self, node: astroid.Module) -> None:
node.inf_type = NoType()
def visit_module(self, node: astroid.Module):
"""
Check for unused XComs.
XComs can be set (pushed) implicitly via return of a python_callable or
execute() of an operator. And explicitly by calling xcom_push().
Currently this only checks unused XComs from return value of a python_callable.
"""
# pylint: disable=too-many-locals,too-many-branches,too-many-nested-blocks
assign_nodes = [n for n in node.body if isinstance(n, astroid.Assign)]
call_nodes = [
n.value for n in assign_nodes if isinstance(n.value, astroid.Call)
]
# Store nodes containing python_callable arg as:
# {task_id: (call node, python_callable func name)}
python_callable_nodes = dict()
for call_node in call_nodes:
if call_node.keywords:
task_id = ""
python_callable = ""
for keyword in call_node.keywords:
if keyword.arg == "python_callable":
python_callable = keyword.value.name
continue
elif keyword.arg == "task_id":
task_id = keyword.value.value
if python_callable:
python_callable_nodes[task_id] = (call_node,
python_callable)
# Now fetch the functions mentioned by python_callable args
xcoms_pushed = dict()
xcoms_pulled_taskids = set()
for (task_id, (python_callable,
callable_func_name)) in python_callable_nodes.items():
if callable_func_name != "<lambda>":
# TODO support lambdas
callable_func = node.getattr(callable_func_name)[0]
if isinstance(callable_func, astroid.FunctionDef):
# Callable_func is str not FunctionDef when imported
callable_func = node.getattr(callable_func_name)[0]
# Check if the function returns any values
if any([
isinstance(n, astroid.Return)
for n in callable_func.body
]):
# Found a return statement
xcoms_pushed[task_id] = (python_callable,
callable_func_name)
# Check if the function pulls any XComs
callable_func_calls = callable_func.nodes_of_class(
astroid.Call)
for callable_func_call in callable_func_calls:
if (isinstance(callable_func_call.func,
astroid.Attribute)
and callable_func_call.func.attrname
== "xcom_pull"):
for keyword in callable_func_call.keywords:
if keyword.arg == "task_ids":
xcoms_pulled_taskids.add(
keyword.value.value)
remainder = xcoms_pushed.keys() - xcoms_pulled_taskids
if remainder:
# There's a remainder in xcoms_pushed_taskids which should've been xcom_pulled.
for remainder_task_id in remainder:
python_callable, callable_func_name = xcoms_pushed[
remainder_task_id]
self.add_message("unused-xcom",
node=python_callable,
args=callable_func_name)
def visit_module(self, node: astroid.Module):
"""Checks in the context of (a) complete DAG(s)."""
dagids_nodes = defaultdict(list)
assigns = node.nodes_of_class(astroid.Assign)
withs = node.nodes_of_class(astroid.With)
def _find_dag(
call_node: astroid.Call, func: Union[astroid.Name,
astroid.Attribute]
) -> Tuple[Union[str, None], Union[astroid.Assign, astroid.Call,
None]]:
"""
Find DAG in a call_node.
:param call_node:
:param func:
:return: (dag_id, node)
:rtype: Tuple
"""
if (hasattr(func, "name")
and func.name == "DAG") or (hasattr(func, "attrname")
and func.attrname == "DAG"):
function_node = safe_infer(func)
if function_node.is_subtype_of(
"airflow.models.DAG") or function_node.is_subtype_of(
"airflow.models.dag.DAG"):
# Check for "dag_id" as keyword arg
if call_node.keywords is not None:
for keyword in call_node.keywords:
# Only constants supported
if keyword.arg == "dag_id" and isinstance(
keyword.value, astroid.Const):
return str(keyword.value.value), call_node
if call_node.args:
if not hasattr(call_node.args[0], "value"):
# TODO Support dynamic dag_id. If dag_id is set from variable, it has no value attr. # pylint: disable=line-too-long
return None, None
return call_node.args[0].value, call_node
return None, None
# Find DAGs in assignments
for assign in assigns:
if isinstance(assign.value, astroid.Call):
func = assign.value.func
dagid, dagnode = _find_dag(assign.value, func)
if dagid and dagnode: # Checks if there are no Nones
dagids_nodes[dagid].append(dagnode)
# Find DAGs in context managers
for with_ in withs:
for with_item in with_.items:
call_node = with_item[0]
if isinstance(call_node, astroid.Call):
func = call_node.func
dagid, dagnode = _find_dag(call_node, func)
if dagid and dagnode: # Checks if there are no Nones
dagids_nodes[dagid].append(dagnode)
# Check if single DAG and if equals filename
# Unit test nodes have file "<?>"
if len(dagids_nodes) == 1 and node.file != "<?>":
dagid, _ = list(dagids_nodes.items())[0]
expected_filename = f"{dagid}.py"
current_filename = node.file.split("/")[-1]
if expected_filename != current_filename:
self.add_message("match-dagid-filename", node=node)
duplicate_dagids = [(dagid, nodes)
for dagid, nodes in dagids_nodes.items()
if len(nodes) >= 2 and dagid is not None]
for (dagid, assign_nodes) in duplicate_dagids:
self.add_message("duplicate-dag-name",
node=assign_nodes[-1],
args=dagid)
