12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697 |
- # Copyright (c) 2016 David Euresti <david@dropbox.com>
-
- """Astroid hooks for typing.py support."""
- import textwrap
-
- from astroid import (
- MANAGER, UseInferenceDefault, extract_node, inference_tip,
- nodes, InferenceError)
- from astroid.nodes import List, Tuple
-
-
- TYPING_NAMEDTUPLE_BASENAMES = {
- 'NamedTuple',
- 'typing.NamedTuple'
- }
-
-
- def infer_typing_namedtuple(node, context=None):
- """Infer a typing.NamedTuple(...) call."""
- # This is essentially a namedtuple with different arguments
- # so we extract the args and infer a named tuple.
- try:
- func = next(node.func.infer())
- except InferenceError:
- raise UseInferenceDefault
-
- if func.qname() != 'typing.NamedTuple':
- raise UseInferenceDefault
-
- if len(node.args) != 2:
- raise UseInferenceDefault
-
- if not isinstance(node.args[1], (List, Tuple)):
- raise UseInferenceDefault
-
- names = []
- for elt in node.args[1].elts:
- if not isinstance(elt, (List, Tuple)):
- raise UseInferenceDefault
- if len(elt.elts) != 2:
- raise UseInferenceDefault
- names.append(elt.elts[0].as_string())
-
- typename = node.args[0].as_string()
- node = extract_node('namedtuple(%(typename)s, (%(fields)s,)) ' %
- {'typename': typename, 'fields': ",".join(names)})
- return node.infer(context=context)
-
-
- def infer_typing_namedtuple_class(node, context=None):
- """Infer a subclass of typing.NamedTuple"""
-
- # Check if it has the corresponding bases
- annassigns_fields = [
- annassign.target.name for annassign in node.body
- if isinstance(annassign, nodes.AnnAssign)
- ]
- code = textwrap.dedent('''
- from collections import namedtuple
- namedtuple({typename!r}, {fields!r})
- ''').format(
- typename=node.name,
- fields=",".join(annassigns_fields)
- )
- node = extract_node(code)
- return node.infer(context=context)
-
-
- def has_namedtuple_base(node):
- """Predicate for class inference tip
-
- :type node: ClassDef
- :rtype: bool
- """
- return set(node.basenames) & TYPING_NAMEDTUPLE_BASENAMES
-
-
- def looks_like_typing_namedtuple(node):
- func = node.func
- if isinstance(func, nodes.Attribute):
- return func.attrname == 'NamedTuple'
- if isinstance(func, nodes.Name):
- return func.name == 'NamedTuple'
- return False
-
-
- MANAGER.register_transform(
- nodes.Call,
- inference_tip(infer_typing_namedtuple),
- looks_like_typing_namedtuple
- )
-
- MANAGER.register_transform(
- nodes.ClassDef,
- inference_tip(infer_typing_namedtuple_class),
- has_namedtuple_base
- )
|