Also add defaults to simple builtin containers

python · JelleZijlstra · Nov 27, 2023 · Jun 2, 2023 · Jun 7, 2023 · Jun 7, 2023
commit 8401f8efb10779f1c04a20213869fac62a7b4736
diff --git a/mypy/stubgen.py b/mypy/stubgen.py
@@ -100,6 +100,7 @@
     NameExpr,
     OpExpr,
     OverloadedFuncDef,
+    SetExpr,
     Statement,
     StrExpr,
     TupleExpr,
@@ -1432,6 +1433,45 @@ def get_str_default_of_node(self, rvalue: Expression) -> str:
             default = repr(rvalue.value)
         elif isinstance(rvalue, BytesExpr):
             default = f"b{rvalue.value!r}"
+        elif isinstance(rvalue, TupleExpr):
+            items_defaults = []
+            for e in rvalue.items:
+                e_default = self.get_str_default_of_node(e)
+                if e_default == "...":
+                    break
+                items_defaults.append(e_default)
+            else:
+                closing = ",)" if len(items_defaults) == 1 else ")"
+                default = "(" + ", ".join(items_defaults) + closing
+        elif isinstance(rvalue, ListExpr):
+            items_defaults = []
+            for e in rvalue.items:
+                e_default = self.get_str_default_of_node(e)
+                if e_default == "...":
+                    break
+                items_defaults.append(e_default)
+            else:
+                default = "[" + ", ".join(items_defaults) + "]"
+        elif isinstance(rvalue, SetExpr):
+            items_defaults = []
+            for e in rvalue.items:
+                e_default = self.get_str_default_of_node(e)
+                if e_default == "...":
+                    break
+                items_defaults.append(e_default)
+            else:
+                if items_defaults:
+                    default = "{" + ", ".join(items_defaults) + "}"
+        elif isinstance(rvalue, DictExpr):
+            items_defaults = []
+            for k, v in rvalue.items:
+                k_default = self.get_str_default_of_node(k)
+                v_default = self.get_str_default_of_node(v)
+                if k_default == "..." or v_default == "...":
+                    break
+                items_defaults.append(f"{k_default}: {v_default}")
+            else:
+                default = "{" + ", ".join(items_defaults) + "}"
 
         if len(default) > 200:  # TODO: what's a good limit?
             default = "..."  # long literals are not useful in stubs

diff --git a/test-data/unit/stubgen.test b/test-data/unit/stubgen.test
@@ -3338,3 +3338,29 @@ def h(x=123456789012345678901234567890123456789012345678901234567890\
 def f(x: str = ...) -> None: ...
 def g(x: bytes = ...) -> None: ...
 def h(x: int = ...) -> None: ...
+
+[case testDefaultsOfBuiltinContainers]
+def f(x=(), y=(1,), z=(1, 2)): ...
+def g(x=[], y=[1, 2]): ...
+def h(x={}, y={1: 2, 3: 4}): ...
+def i(x={1, 2, 3}): ...
+def j(x=[(1,"a"), (2,"b")]): ...
+
+[out]
+def f(x=(), y=(1,), z=(1, 2)) -> None: ...
+def g(x=[], y=[1, 2]) -> None: ...
+def h(x={}, y={1: 2, 3: 4}) -> None: ...
+def i(x={1, 2, 3}) -> None: ...
+def j(x=[(1, 'a'), (2, 'b')]) -> None: ...
+
+[case testDefaultsOfBuiltinContainersWithNonTrivialContent]
+def f(x=(1, u.v), y=(k(),), z=(w,)): ...
+def g(x=[1, u.v], y=[k()], z=[w]): ...
+def h(x={1: u.v}, y={k(): 2}, z={w: w}): ...
+def i(x={u.v, 2}, y={3, k()}, z={w}): ...
+
+[out]
+def f(x=..., y=..., z=...) -> None: ...
+def g(x=..., y=..., z=...) -> None: ...
+def h(x=..., y=..., z=...) -> None: ...
+def i(x=..., y=..., z=...) -> None: ...