summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorcpopa <devnull@localhost>2014-08-13 13:52:45 +0300
committercpopa <devnull@localhost>2014-08-13 13:52:45 +0300
commita6c0d3d012ce9624546f2d22aa57e5e8ed3fcd69 (patch)
tree9a390bdcd3813117d5fafe24a7bc4c15cad61bee
parent6e2922b562c7beda407a2cb8e4f00d6d0e1db682 (diff)
downloadastroid-a6c0d3d012ce9624546f2d22aa57e5e8ed3fcd69.tar.gz
Support iterable of tuples as well as iterable of names.
-rw-r--r--brain/py2stdlib.py16
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: