diff options
author | cpopa <devnull@localhost> | 2014-08-13 13:52:45 +0300 |
---|---|---|
committer | cpopa <devnull@localhost> | 2014-08-13 13:52:45 +0300 |
commit | a6c0d3d012ce9624546f2d22aa57e5e8ed3fcd69 (patch) | |
tree | 9a390bdcd3813117d5fafe24a7bc4c15cad61bee | |
parent | 6e2922b562c7beda407a2cb8e4f00d6d0e1db682 (diff) | |
download | astroid-a6c0d3d012ce9624546f2d22aa57e5e8ed3fcd69.tar.gz |
Support iterable of tuples as well as iterable of names.
-rw-r--r-- | brain/py2stdlib.py | 16 |
1 files changed, 12 insertions, 4 deletions
diff --git a/brain/py2stdlib.py b/brain/py2stdlib.py index 2b879f0..05ce76f 100644 --- a/brain/py2stdlib.py +++ b/brain/py2stdlib.py @@ -52,15 +52,23 @@ def infer_func_form(node, base_type, context=None, enum=False): else: # Enums supports either iterator of (name, value) pairs # or mappings. - # TODO: support only list and mappings. + # TODO: support only list, tuples and mappings. if hasattr(names, 'items') and isinstance(names.items, list): attributes = [infer_first(const[0]).value for const in names.items if isinstance(const[0], nodes.Const)] elif hasattr(names, 'elts'): - attributes = [infer_first(const.elts[0]).value - for const in names.elts - if isinstance(const, nodes.Tuple)] + # Enums can support either ["a", "b", "c"] + # or [("a", 1), ("b", 2), ...], but they can't + # be mixed. + if all(isinstance(const, nodes.Tuple) + for const in names.elts): + attributes = [infer_first(const.elts[0]).value + for const in names.elts + if isinstance(const, nodes.Tuple)] + else: + attributes = [infer_first(const).value + for const in names.elts] else: raise AttributeError if not attributes: |