Fix an error with enums, fixes #1675

This commit is contained in:
Dave Halter
2020-10-23 20:25:00 +02:00
parent a5a36a049c
commit e671a0cb6d
2 changed files with 27 additions and 1 deletions

View File

@@ -819,7 +819,8 @@ def get_metaclass_filters(func):
and metaclass.get_root_context().py__name__() == 'enum':
filter_ = ParserTreeFilter(parent_context=cls.as_context())
return [DictFilter({
name.string_name: EnumInstance(cls, name).name for name in filter_.values()
name.string_name: EnumInstance(cls, name).name
for name in filter_.values()
})]
return func(cls, metaclasses, is_instance)
return wrapper
@@ -837,6 +838,14 @@ class EnumInstance(LazyValueWrapper):
return ValueName(self, self._name.tree_name)
def _get_wrapped_value(self):
n = self._name.string_name
if n.startswith('__') and n.endswith('__') or self._name.api_type == 'function':
inferred = self._name.infer()
if inferred:
return next(iter(inferred))
o, = self.inference_state.builtins_module.py__getattribute__('object')
return o
value, = self._cls.execute_with_values()
return value

View File

@@ -340,3 +340,20 @@ def test_overload(Script, code):
x1, x2 = Script(code, path=os.path.join(dir_, 'foo.py')).get_signatures()
assert x1.to_string() == 'with_overload(x: int, y: int) -> float'
assert x2.to_string() == 'with_overload(x: str, y: list) -> float'
def test_enum(Script):
script = Script('''\
from enum import Enum
class Planet(Enum):
MERCURY = (3.303e+23, 2.4397e6)
VENUS = (4.869e+24, 6.0518e6)
def __init__(self, mass, radius):
self.mass = mass # in kilograms
self.radius = radius # in meters
Planet.MERCURY''')
completion, = script.complete()
assert not completion.get_signatures()