From 40e1e3f560fdb9d7afd85d9337cf767158f57c0b Mon Sep 17 00:00:00 2001
From: Alex Ford <alex.ford@abcellera.com>
Date: Mon, 10 Oct 2022 11:15:27 -0700
Subject: [PATCH] Extend dataclass constructor hinting to attrs next-gen apis.

Trivially extends dataclass constructor hinting to attrs next-gen APIs.

This will stumble in cases where attrs extends beyond the standard
dataclasses API, such as complex use of defaults, converters, et al.
However, it likely covers the vast majority of cases which fall solidly
in the intersection of the two APIs.

Extension beyond these cases could use [PEP0681 dataclass_transforms],
however this is definitely a problem for another day.

[PEP0681 dataclass_transforms]: https://peps.python.org/pep-0681/

https://github.com/davidhalter/jedi/issues/1835
---
 jedi/plugins/stdlib.py                |  9 +++++++
 test/test_inference/test_signature.py | 39 +++++++++++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/jedi/plugins/stdlib.py b/jedi/plugins/stdlib.py
index 17f1df3b..81803090 100644
--- a/jedi/plugins/stdlib.py
+++ b/jedi/plugins/stdlib.py
@@ -803,6 +803,15 @@ _implemented = {
         # For now this works at least better than Jedi trying to understand it.
         'dataclass': _dataclass
     },
+    # attrs exposes declaration interface roughly compatible with dataclasses
+    # via attrs.define, attrs.frozen and attrs.mutable
+    # https://www.attrs.org/en/stable/names.html
+    'attr': {
+        'define': _dataclass,
+    },
+    'attrs': {
+        'define': _dataclass,
+    },
     'os.path': {
         'dirname': _create_string_input_function(os.path.dirname),
         'abspath': _create_string_input_function(os.path.abspath),
diff --git a/test/test_inference/test_signature.py b/test/test_inference/test_signature.py
index 9ee8a55c..b471a19d 100644
--- a/test/test_inference/test_signature.py
+++ b/test/test_inference/test_signature.py
@@ -355,6 +355,45 @@ def test_dataclass_signature(Script, skip_pre_python37, start, start_params):
     price, = sig.params[-2].infer()
     assert price.name == 'float'
 
+@pytest.mark.parametrize(
+    'start, start_params', [
+        ['@define\nclass X:', []],
+        ['@frozen\nclass X:', []],
+        ['@mutable\nclass X:', []],
+        ['@define(eq=True)\nclass X:', []],
+        [dedent('''
+         class Y():
+             y: int
+         @define
+         class X(Y):'''), []],
+        [dedent('''
+         @define
+         class Y():
+             y: int
+             z = 5
+         @define
+         class X(Y):'''), ['y']],
+    ]
+)
+def test_attrs_signature(Script, skip_pre_python37, start, start_params):
+    code = dedent('''
+            name: str
+            foo = 3
+            price: float
+            quantity: int = 0.0
+
+        X(''')
+
+    # attrs exposes two namespaces
+    code = 'from attrs import define, frozen\n' + start + code
+
+    sig, = Script(code).get_signatures()
+    assert [p.name for p in sig.params] == start_params + ['name', 'price', 'quantity']
+    quantity, = sig.params[-1].infer()
+    assert quantity.name == 'int'
+    price, = sig.params[-2].infer()
+    assert price.name == 'float'
+
 
 @pytest.mark.parametrize(
     'stmt, expected', [