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 _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):
"""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)