Support default_factory for structs

example/app/interfaces/dates.pyi

@@ -13,3 +13,7 @@ class DateTime:
 
 @dataclass
 class Date: ...
+
+EPOCH: DateTime = DateTime(
+    year=1970, month=1, day=1, hour=0, minute=0, second=0, microsecond=0
+)

example/app/interfaces/shared.pyi

@@ -2,7 +2,9 @@ from dataclasses import dataclass
 from typing import *
 
 class NotFound(Exception):
-    def __init__(self, message: Optional[str] = "Not Found"): ...
+    message: Optional[str] = "Not Found"
+
+    def __init__(self, message: Optional[str] = "Not Found") -> None: ...
 
 class EmptyException(Exception): ...
 
@@ -11,5 +13,9 @@ class LimitOffset:
     limit: Optional[int] = None
     offset: Optional[int] = None
 
+INT_CONST_1: int = 1234
+MAP_CONST: Dict[str, str] = {"hello": "world", "goodnight": "moon"}
+INT_CONST_2: int = 1234
+
 class Service:
     def ping(self) -> str: ...

example/app/interfaces/todo.pyi

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from enum import IntEnum
 from typing import *
 from . import dates
@@ -13,19 +13,26 @@ class TodoType(IntEnum):
 class TodoItem:
     id: int
     text: str
-    type: int
+    type: TodoType
     created: dates.DateTime
     is_deleted: bool
-    picture: Optional[str] = None
+    picture: Optional[bytes] = None
     is_favorite: bool = False
 
+@dataclass
+class TodoCounter:
+    todos: Dict[int, TodoItem] = field(default_factory=dict)
+    plain_ids: Set[int] = field(default_factory=lambda: {1, 2, 3})
+    note_ids: List[int] = field(default_factory=list)
+    checkboxes_ids: Set[int] = field(default_factory=set)
+
 class Todo:
-    def create(self, text: str, type: int) -> int: ...
+    def create(self, text: str, type: TodoType) -> int: ...
     def update(self, item: TodoItem) -> None: ...
     def get(self, id: int) -> TodoItem: ...
     def all(self, pager: shared.LimitOffset) -> List[TodoItem]: ...
     def filter(self, ids: List[int]) -> List[TodoItem]: ...
     def stats(self) -> Dict[int, float]: ...
     def types(self) -> Set[int]: ...
-    def groupby(self) -> Dict[int, List[TodoItem]]: ...
+    def groupby(self) -> Dict[TodoType, List[TodoItem]]: ...
     def ping(self) -> str: ...

example/app/server.py

@@ -5,13 +5,13 @@
 from thriftpy2.rpc import make_server
 
 if TYPE_CHECKING:
-    from example.app.interfaces.todo import TodoItem
+    from example.app.interfaces.todo import TodoItem, TodoType
 
 todos: Dict[int, "TodoItem"] = {}
 
 
 class Dispatcher:
-    def create(self, text: str, type: int) -> int:
+    def create(self, text: str, type: TodoType) -> int:
         todo_id = max(todos.keys() or [0]) + 1
         created = datetime.datetime.now()
         todos[todo_id] = interfaces.todo.TodoItem(

example/interfaces/todo.thrift

@@ -21,9 +21,18 @@ struct TodoItem {
     7: required bool is_favorite = false
 }
 
+
 typedef list<TodoItem> TodoItemList
 
 
+struct TodoCounter {
+    1: required map<i32, TodoItem> todos = {}
+    2: required set<i32> plain_ids = [1, 2, 3]
+    3: required list<i32> note_ids = []
+    4: required set<i32> checkboxes_ids = []
+}
+
+
 service Todo extends shared.Service {
     i32 create(
         1: string text,

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "thrift-pyi"
-version = "1.0.0"
+version = "1.1.0"
 description = "This is simple `.pyi` stubs generator from thrift interfaces"
 readme = "README.rst"
 repository = "https://github.com/unmade/thrift-pyi"

src/thriftpyi/entities.py

@@ -2,7 +2,15 @@
 
 import ast
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Sequence, Type, Union
+from typing import (
+    TYPE_CHECKING,
+    MutableMapping,
+    MutableSequence,
+    MutableSet,
+    Sequence,
+    Type,
+    Union,
+)
 
 if TYPE_CHECKING:
     AnyFunctionDef = Union[ast.AsyncFunctionDef, ast.FunctionDef]
@@ -161,7 +169,7 @@ def as_ast(self) -> ast.AnnAssign | ast.Assign:
         if self.type is None:
             return ast.Assign(
                 targets=[ast.Name(id=self.name, ctx=ast.Store())],
-                value=ast.Constant(value=self.value, kind=None),
+                value=self._make_ast_value(),
                 lineno=0,
             )
 
@@ -173,11 +181,38 @@ def as_ast(self) -> ast.AnnAssign | ast.Assign:
         if self.required and self.value is None:
             value = None
         else:
-            value = ast.Constant(value=self.value, kind=None)
+            value = self._make_ast_value()
 
         return ast.AnnAssign(
             target=ast.Name(id=self.name, ctx=ast.Store()),
             annotation=annotation,
             value=value,
             simple=1,
         )
+
+    def _make_ast_value(self) -> ast.expr:
+        return ast.Constant(value=self.value, kind=None)
+
+
+@dataclass
+class StructField(Field):
+    def _make_ast_value(self) -> ast.expr:
+        if isinstance(self.value, (MutableSequence, MutableSet, MutableMapping)):
+            if self.value:
+                value = ast.Lambda(
+                    args=[], body=ast.Constant(value=self.value, kind=None)
+                )
+            else:
+                value = ast.Name(id=self.value.__class__.__name__, ctx=ast.Load())
+
+            return ast.Call(
+                func=ast.Name(id="field", ctx=ast.Load()),
+                args=[],
+                keywords=[
+                    ast.keyword(
+                        arg="default_factory",
+                        value=value,
+                    )
+                ],
+            )
+        return super()._make_ast_value()

src/thriftpyi/proxies.py

@@ -2,7 +2,7 @@
 
 from typing import cast
 
-from thriftpyi.entities import Field, FieldValue, Method, ModuleItem
+from thriftpyi.entities import Field, FieldValue, Method, ModuleItem, StructField
 from thriftpyi.utils import get_python_type, guess_type
 
 
@@ -127,7 +127,7 @@ def _make_service(tservice) -> ModuleItem:
 
     @staticmethod
     def _make_struct(tclass) -> ModuleItem:
-        spec = TSpecProxy(
+        spec = TStructSpecProxy(
             module_name=tclass.__module__,
             thrift_spec=tclass.thrift_spec,
             default_spec=dict(tclass.default_spec),
@@ -188,3 +188,16 @@ def _get_python_type(self, item: TSpecItemProxy) -> str:
     def _get_default_value(self, item: TSpecItemProxy) -> FieldValue:
         default_value = self.default_spec.get(item.name)
         return cast(FieldValue, default_value)
+
+
+class TStructSpecProxy(TSpecProxy):
+    def get_fields(self, *, ignore_type: bool = False) -> list[Field]:
+        return [
+            StructField(
+                name=item.name,
+                type=self._get_python_type(item) if not ignore_type else None,
+                value=self._get_default_value(item),
+                required=item.required,
+            )
+            for item in self.thrift_spec
+        ]

src/thriftpyi/stubs.py

@@ -33,7 +33,7 @@ def build_init(imports: Iterable[str]) -> ast.Module:
 def _make_imports(proxy: TModuleProxy) -> list[ast.ImportFrom]:
     imports = []
     if proxy.has_structs():
-        imports.append(_make_absolute_import("dataclasses", "dataclass"))
+        imports.append(_make_absolute_import("dataclasses", "dataclass, field"))
     if proxy.has_enums():
         imports.append(_make_absolute_import("enum", "IntEnum"))
 

tests/stubs/expected/async/dates.pyi

@@ -13,3 +13,7 @@ class DateTime:
 
 @dataclass
 class Date: ...
+
+EPOCH: DateTime = DateTime(
+    year=1970, month=1, day=1, hour=0, minute=0, second=0, microsecond=0
+)

tests/stubs/expected/async/shared.pyi

@@ -1,10 +1,6 @@
 from dataclasses import dataclass
 from typing import *
 
-INT_CONST_1: int = 1234
-MAP_CONST: Dict[str, str] = {"hello": "world", "goodnight": "moon"}
-INT_CONST_2: int = 1234
-
 class NotFound(Exception):
     message: Optional[str] = "Not Found"
 
@@ -17,5 +13,9 @@ class LimitOffset:
     limit: Optional[int] = None
     offset: Optional[int] = None
 
+INT_CONST_1: int = 1234
+MAP_CONST: Dict[str, str] = {"hello": "world", "goodnight": "moon"}
+INT_CONST_2: int = 1234
+
 class Service:
     async def ping(self) -> str: ...

tests/stubs/expected/async/todo.pyi

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from enum import IntEnum
 from typing import *
 from . import dates
@@ -19,6 +19,13 @@ class TodoItem:
     picture: Optional[bytes] = None
     is_favorite: bool = False
 
+@dataclass
+class TodoCounter:
+    todos: Dict[int, TodoItem] = field(default_factory=dict)
+    plain_ids: Set[int] = field(default_factory=lambda: {1, 2, 3})
+    note_ids: List[int] = field(default_factory=list)
+    checkboxes_ids: Set[int] = field(default_factory=set)
+
 class Todo:
     async def create(self, text: str, type: TodoType) -> int: ...
     async def update(self, item: TodoItem) -> None: ...

tests/stubs/expected/optional/shared.pyi

@@ -1,10 +1,6 @@
 from dataclasses import dataclass
 from typing import *
 
-INT_CONST_1: int = 1234
-MAP_CONST: Dict[str, str] = {"hello": "world", "goodnight": "moon"}
-INT_CONST_2: int = 1234
-
 class NotFound(Exception):
     message: Optional[str] = "Not Found"
 
@@ -17,5 +13,9 @@ class LimitOffset:
     limit: Optional[int] = None
     offset: Optional[int] = None
 
+INT_CONST_1: int = 1234
+MAP_CONST: Dict[str, str] = {"hello": "world", "goodnight": "moon"}
+INT_CONST_2: int = 1234
+
 class Service:
     def ping(self) -> str: ...

tests/stubs/expected/optional/todo.pyi

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from enum import IntEnum
 from typing import *
 from . import dates
@@ -19,6 +19,13 @@ class TodoItem:
     picture: Optional[bytes] = None
     is_favorite: Optional[bool] = False
 
+@dataclass
+class TodoCounter:
+    todos: Optional[Dict[int, TodoItem]] = field(default_factory=dict)
+    plain_ids: Optional[Set[int]] = field(default_factory=lambda: {1, 2, 3})
+    note_ids: Optional[List[int]] = field(default_factory=list)
+    checkboxes_ids: Optional[Set[int]] = field(default_factory=set)
+
 class Todo:
     def create(self, text: str, type: TodoType) -> int: ...
     def update(self, item: TodoItem) -> None: ...

tests/stubs/expected/sync/todo.pyi

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from enum import IntEnum
 from typing import *
 from . import dates
@@ -19,6 +19,13 @@ class TodoItem:
     picture: Optional[bytes] = None
     is_favorite: bool = False
 
+@dataclass
+class TodoCounter:
+    todos: Dict[int, TodoItem] = field(default_factory=dict)
+    plain_ids: Set[int] = field(default_factory=lambda: {1, 2, 3})
+    note_ids: List[int] = field(default_factory=list)
+    checkboxes_ids: Set[int] = field(default_factory=set)
+
 class Todo:
     def create(self, text: str, type: TodoType) -> int: ...
     def update(self, item: TodoItem) -> None: ...

tests/test_cli.py

@@ -10,8 +10,8 @@
     "expected_dir,args",
     [
         ("tests/stubs/expected/sync", ["--strict-optional"]),
-        # ("tests/stubs/expected/async", ["--async", "--strict-optional"]),
-        # ("tests/stubs/expected/optional", []),
+        ("tests/stubs/expected/async", ["--async", "--strict-optional"]),
+        ("tests/stubs/expected/optional", []),
     ],
 )
 def test_main(capsys, expected_dir, args):