From 93c52f615ae9fe1ee4481b7f36a834e51e4745f9 Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Fri, 19 Jul 2019 11:35:13 +0200
Subject: [PATCH] Get inheritance of dataclass right

---
 CHANGELOG.rst                        |  2 +-
 jedi/plugins/stdlib.py               |  7 ++++---
 test/test_evaluate/test_signature.py | 28 ++++++++++++++++++++++------
 3 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index db5f016f..5a6a3939 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -6,7 +6,7 @@ Changelog
 0.14.2 (XXXX-XX-XX)
 +++++++++++++++++++
 
-- Better support for enum completions
+- Better support for enums/dataclasses
 
 
 0.14.1 (2019-07-13)
diff --git a/jedi/plugins/stdlib.py b/jedi/plugins/stdlib.py
index db09896c..a87fb4ed 100644
--- a/jedi/plugins/stdlib.py
+++ b/jedi/plugins/stdlib.py
@@ -24,6 +24,7 @@ from jedi.evaluate.base_context import ContextualizedNode, \
     NO_CONTEXTS, ContextSet, ContextWrapper, LazyContextWrapper
 from jedi.evaluate.context import ClassContext, ModuleContext, \
     FunctionExecutionContext
+from jedi.evaluate.context.klass import ClassMixin
 from jedi.evaluate.context import iterable
 from jedi.evaluate.lazy_context import LazyTreeContext, LazyKnownContext, \
     LazyKnownContexts
@@ -530,11 +531,11 @@ def _dataclass(obj, arguments):
     return NO_CONTEXTS
 
 
-class DataclassWrapper(ContextWrapper):
+class DataclassWrapper(ContextWrapper, ClassMixin):
     def get_signatures(self):
         param_names = []
-        for cls in reversed(list(self._wrapped_context.py__mro__())):
-            if isinstance(cls, ClassContext) and not cls.is_stub():
+        for cls in reversed(list(self.py__mro__())):
+            if isinstance(cls, DataclassWrapper):
                 filter_ = cls.get_global_filter()
                 for name in filter_.values():
                     d = name.tree_name.get_definition()
diff --git a/test/test_evaluate/test_signature.py b/test/test_evaluate/test_signature.py
index 4a97e2e9..9a95fd7c 100644
--- a/test/test_evaluate/test_signature.py
+++ b/test/test_evaluate/test_signature.py
@@ -64,18 +64,34 @@ def test_pow_signature(Script):
                        'pow(x: int, y: int, /) -> Any'}
 
 
-@pytest.mark.parametrize('decorator', ['@dataclass', '@dataclass(eq=True)'])
-def test_dataclass_signature(Script, skip_pre_python37, decorator):
+@pytest.mark.parametrize(
+    'start, start_params', [
+        ['@dataclass\nclass X:', []],
+        ['@dataclass(eq=True)\nclass X:', []],
+        [dedent('''
+         class Y():
+             y: int
+         @dataclass
+         class X(Y):'''), []],
+        [dedent('''
+         @dataclass
+         class Y():
+             y: int
+             z = 5
+         @dataclass
+         class X(Y):'''), ['y']],
+    ]
+)
+def test_dataclass_signature(Script, skip_pre_python37, start, start_params):
     code = dedent('''
-        class InventoryItem:
             name: str
             foo = 3
             unit_price: float
             quantity_on_hand: int = 0
 
-        InventoryItem(''')
+        X(''')
 
-    code = 'from dataclasses import dataclass\n' + decorator + code
+    code = 'from dataclasses import dataclass\n' + start + code
 
     sig, = Script(code).call_signatures()
-    assert [p.name for p in sig.params] == ['name', 'unit_price', 'quantity_on_hand']
+    assert [p.name for p in sig.params] == start_params + ['name', 'unit_price', 'quantity_on_hand']