Repository: unmade/thrift-pyi

Issue:
Pass through underlying Thrift type as field metadata
Hello again!

One issue we are encountering is that although `int` is the valid Python type for both `i32` and `i64`, we have some use cases that require access to the underlying Thrift type. One solution that feels ergonomic enough would be if the generated classes looked something like:

```python
from dataclasses import dataclass, field

@dataclass
class User:
    id: int = field(metadata={'thrift_type': 'i32'})
    username: str
    last_login: int = field(metadata={'thrift_type': 'i64'})
    score: float
    is_admin: bool
```

The annotations can be on all types (or maybe all primitive types? or maybe including structs and just serialize the struct name to string?) or just to disambiguate `int`. I don't really have a preference for right now since `int` is the most pressing use case for me.

My understanding is that the original intent of this library was to generate `.pyi` type stubs for Thrift interfaces, and that this change might push things closer to a "Thrift to Dataclass transpiler". Are you open to this kind of change? 

I think I see a path for how I could make this change myself, but would appreciate some tips from you.

Hints from issue comments:
@FrankPortman: I played around with this a little bit today and I have a sketch of a solution that looks like it works. It does involve decoupling `StructField` from `Field` in terms of class hierarchy but if we go down this route, I think that makes sense. They would then be fairly diverging concepts.

Example of code it generates:

```
@dataclass
class TodoItem:
    id: int = field(metadata={"thrift_type": "I32"})
    text: str = field(metadata={"thrift_type": "STRING"})
    type: TodoType = field(metadata={"thrift_type": "I32"})
    created: dates.DateTime = field(metadata={"thrift_type": "STRUCT"})
    is_deleted: bool = field(metadata={"thrift_type": "BOOL"})
    picture: Optional[bytes] = field(default=None, metadata={"thrift_type": "BINARY"})
    is_favorite: bool = field(default=False, metadata={"thrift_type": "BOOL"})
```

The string is generated by https://github.com/Thriftpy/thriftpy2/blob/master/thriftpy2/thrift.py#L30 but could just pass through the `thriftpy2` `Enum` as well.@unmade: Hey Frank,

That surely possible and I'll give it a look until the EOW. Feel free to open a PR with any code samples you have as it might be useful.

I think with changes like these it is OK that `StructField` is decoupled from Field, although it would be better if `StructField` is still complies to `Field` interface. 

Just to so I have a better understanding - are you planning to retrieve that information with some other mypy/pyright plugin? Is it possible to retrieve it from `.pyi` files? I don't know much about what plugins can and can't do@FrankPortman: Sounds great - I will send you a rough sketch of a solution soon.

> Just to so I have a better understanding - are you planning to retrieve that information with some other mypy/pyright plugin? Is it possible to retrieve it from .pyi files? I don't know much about what plugins can and can't do

One of the reasons I felt you may not be open to this change, is that we actually use `thrift-pyi` quite differently than its stated purpose. We do use Thrift heavily, but as a place to hold cross-language schemas and constants. And so `thrift-pyi` ends up being a very nice simple way to codegen constants and data-classes, whereas the Apache Thrift compiler is focused more on the SerDe aspect. We basically use `thrift-pyi` as a binary, and do some basic munging after such as renaming `.pyi` files to `.py`, etc.

Very specifically, we just need some way to retain the "real" type that the developer intended in Thrift, because Python has no distinction between `i32` and `i64` (but our other languages do). We use this for things such as deriving Spark schemas, which DO have a distinction between `i32` and `i64`, etc. Maintaining this behavior for `STRUCT`, `BOOL`, etc. I just added into my example for completeness, but this is really only impacting us for `int`.

It's great that you're open to this, but if you think about it more and feel it is a departure from the original goal of `thrift-pyi` then we would likely just fork the parts we need (and can keep it OSS as well).@unmade: Wow, thanks for such detailed explanation! 

I think I can work on adding that thrift metadata, but then with your example it will only be present for dataclass fields. 

What do you think if instead we use [typing.Annotated](https://docs.python.org/3/library/typing.html) to support it everywhere? Something like this:

```python
i32 = Annotated[int, {"thrift_type": "I32"}]

@dataclass
class TodoItem:
    id: i32

class Todo:
    def get(self, id: i32) -> TodoItem:
```

That is just the concept, but looks like it should be doable. Let me know what you think@FrankPortman: I think that would work - my understanding is we can access that `thrift_type` metadata value at runtime for our purposes. To be sure, let me check out a dummy example though.

That would also allow you to generate a handful of these aliases at the project level, and then not clutter up the actual generated code. And also as far as I can tell there would be no changes to anything but the type annotation itself, so it's probably a smaller code change than what I have on my branch. 

I think this should work - but to be clear on our end we don't have any specific asks to support this in regular fields (only the dataclass defs). So it's really up to you if you want this annotation everywhere.@unmade: Got it. To me using `Annotated` would be preferable way of implementing this, so waiting for your confirmation that it is going to work for you as well@FrankPortman: Alright this looks like it will work. We will have to handle some parsing to make sure we detect `I32`, etc inside `LIST<..>` or `MAP<...>` but that is true in any solution.

I think this helper inside `thriftpy2` basically gives you what you want as far as mappings: https://github.com/Thriftpy/thriftpy2/blob/master/thriftpy2/thrift.py#L29

Thank you!@FrankPortman: What I used to test:

```python
from dataclasses import dataclass, fields
from typing import Annotated, get_type_hints, get_origin, get_args

# Define annotated types for clarity.
i32 = Annotated[int, {"thrift_type": "I32"}]
string = Annotated[str, {"thrift_type": "STRING"}]
list_of_strings = Annotated[list[str], {"thrift_type": "LIST<STRING>"}]
dict_str_int = Annotated[dict[str, int], {"thrift_type": "MAP<STRING, I32>"}]

@dataclass
class TodoItem:
    id: i32
    title: string
    tags: list_of_strings
    attributes: dict_str_int

def inspect_dataclass(cls):
    print(f"Inspecting {cls.__name__}:")
    hints = get_type_hints(cls, include_extras=True)
    for field in fields(cls):
        name = field.name
        ann = hints[name]
        print(f"\nField: {name}")
        if get_origin(ann) is Annotated:
            base_type, *meta = get_args(ann)
            print(f"  Base type: {base_type}")
            print(f"  Metadata: {meta}")
        else:
            print(f"  Type: {ann}")


inspect_dataclass(TodoItem)
```





@unmade: @FrankPortman 
I opened PR with requested changes. Would you be able to test from the branch. I believe you can install the library like this:
```python
pip install git+https://github.com/unmade/thrift-pyi@56-meta-info-about-original-thrift-type
```

Let me know if it works for you or you expected something else. If everything is OK then I'll merge the changes@FrankPortman: Thank you! We will check tomorrow or Thursday at the latest.@FrankPortman: Update -

I am still testing it. I hit some hiccups with our build system but everything looks wired up correctly now.

I am seeing a lot of errors like this:

```python
from dataclasses import dataclass
from enum import IntEnum
from typing import *
from . import _typedefs

class EnvironmentType(IntEnum):
    Prod = 0
    Dev = 1

@dataclass
class Environment:
    environmentType: EnvironmentType
    gitSha: Optional[_typedefs.String] = None
```

leading to an error like this

```
    gitSha: Optional[_typedefs.String] = None
  File "/lib/python3.10/typing.py", line 312, in inner
    return func(*args, **kwds)
  File "/lib/python3.10/typing.py", line 403, in __getitem__
    return self._getitem(self, parameters)
  File "/lib/python3.10/typing.py", line 530, in Optional
    return Union[arg, type(None)]
  File "/lib/python3.10/typing.py", line 312, in inner
    return func(*args, **kwds)
  File "/lib/python3.10/typing.py", line 403, in __getitem__
    return self._getitem(self, parameters)
  File "/lib/python3.10/typing.py", line 516, in Union
    parameters = _remove_dups_flatten(parameters)
  File "/lib/python3.10/typing.py", line 281, in _remove_dups_flatten
    return tuple(_deduplicate(params))
  File "/lib/python3.10/typing.py", line 255, in _deduplicate
    all_params = set(params)
  File "/lib/python3.10/typing.py", line 1651, in __hash__
    return hash((self.__origin__, self.__metadata__))
TypeError: unhashable type: 'dict'
```

It seems like it is unhappy with the metadata being a dict? Perhaps we can use a frozen dataclass to hold the metadata? 
Realistically the metadata is entirely controlled by this library, so I also think it would okay if it was simply a `str` until an n=2 case arose. 

Creating the frozen dataclass metadata holder would be more futureproof though if you decide to add more annotations, such as eventually passing through field annotations (e.g. https://github.com/apache/thrift/blob/master/test/AnnotationTest.thrift).
@unmade: Thanks for testing it! It is very strange that I didn't catch that error in testing

In all cases I think it is better to store metadata as either a frozen dataclass. I pushed the changes, so please check it :) 

@FrankPortman: ```python
from dataclasses import dataclass
from typing import Annotated

@dataclass()
class Metadata(frozen=True):
    thrift_type: str

Binary = Annotated[bytes, Metadata(thrift_type="BINARY")]
Bool = Annotated[bool, Metadata(thrift_type="BOOL")]
Byte = Annotated[int, Metadata(thrift_type="BYTE")]
Double = Annotated[float, Metadata(thrift_type="DOUBLE")]
I16 = Annotated[int, Metadata(thrift_type="I16")]
I32 = Annotated[int, Metadata(thrift_type="I32")]
I64 = Annotated[int, Metadata(thrift_type="I64")]
String = Annotated[str, Metadata(thrift_type="STRING")]
```

That is the code it generates - looks like `frozen=True` is not going into `dataclass` args which is causing `TypeError: Metadata.__init_subclass__() takes no keyword arguments` upon import.

See this comment: https://github.com/unmade/thrift-pyi/pull/57/files#r1960004116@unmade: Whoops, indeed! I had it right in the test file, but the file itself was not included in the comparison, so the mistake wasn't catch.

I pushed the fix, sorry for the inconvenience! @FrankPortman: Okay I am done testing 😅 - I am relatively new to doing fancy reflection things in Python so I wanted to be sure it was comprehensive.

Everything looks good to me, and I can submit a PR for any minor fixes if things come up moving forward.

Just a note, you may want to be clear in release notes that this change may be breaking if people downstream have naive uses of `get_origin` and `get_args` and whatnot.

Thanks again for this!@unmade: Perfect, thank you! 

Please check it out in version `2.0.0`

Solution patch:
diff --git a/example/app/interfaces/_typedefs.pyi b/example/app/interfaces/_typedefs.pyi
new file mode 100644
index 0000000..f4aadd6
--- /dev/null
+++ b/example/app/interfaces/_typedefs.pyi
@@ -0,0 +1,10 @@
+from typing import Annotated
+
+Binary = Annotated[bytes, {"thrift_type": "BINARY"}]
+Bool = Annotated[bool, {"thrift_type": "BOOL"}]
+Byte = Annotated[int, {"thrift_type": "BYTE"}]
+Double = Annotated[float, {"thrift_type": "DOUBLE"}]
+I16 = Annotated[int, {"thrift_type": "I16"}]
+I32 = Annotated[int, {"thrift_type": "I32"}]
+I64 = Annotated[int, {"thrift_type": "I64"}]
+String = Annotated[str, {"thrift_type": "STRING"}]
diff --git a/example/app/interfaces/dates.pyi b/example/app/interfaces/dates.pyi
index 83e18c2..c0a1f51 100644
--- a/example/app/interfaces/dates.pyi
+++ b/example/app/interfaces/dates.pyi
@@ -1,15 +1,16 @@
 from dataclasses import dataclass
 from typing import *
+from . import _typedefs
 
 @dataclass
 class DateTime:
-    year: int
-    month: int
-    day: int
-    hour: int
-    minute: int
-    second: int
-    microsecond: Optional[int] = None
+    year: _typedefs.I16
+    month: _typedefs.Byte
+    day: _typedefs.Byte
+    hour: _typedefs.I16
+    minute: _typedefs.Byte
+    second: _typedefs.Byte
+    microsecond: Optional[_typedefs.I64] = None
 
 @dataclass
 class Date: ...
diff --git a/example/app/interfaces/shared.pyi b/example/app/interfaces/shared.pyi
index ff9725a..f6bbd56 100644
--- a/example/app/interfaces/shared.pyi
+++ b/example/app/interfaces/shared.pyi
@@ -1,21 +1,22 @@
 from dataclasses import dataclass
 from typing import *
+from . import _typedefs
 
 class NotFound(Exception):
-    message: Optional[str] = "Not Found"
+    message: Optional[_typedefs.String] = "Not Found"
 
-    def __init__(self, message: Optional[str] = "Not Found") -> None: ...
+    def __init__(self, message: Optional[_typedefs.String] = "Not Found") -> None: ...
 
 class EmptyException(Exception): ...
 
 @dataclass
 class LimitOffset:
-    limit: Optional[int] = None
-    offset: Optional[int] = None
+    limit: Optional[_typedefs.I32] = None
+    offset: Optional[_typedefs.I32] = 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: ...
+    def ping(self) -> _typedefs.String: ...
diff --git a/example/app/interfaces/todo.pyi b/example/app/interfaces/todo.pyi
index 27bc446..f520042 100644
--- a/example/app/interfaces/todo.pyi
+++ b/example/app/interfaces/todo.pyi
@@ -1,6 +1,7 @@
 from dataclasses import dataclass, field
 from enum import IntEnum
 from typing import *
+from . import _typedefs
 from . import dates
 from . import shared
 
@@ -11,32 +12,32 @@ class TodoType(IntEnum):
 
 @dataclass
 class TodoItem:
-    id: int
-    text: str
+    id: _typedefs.I32
+    text: _typedefs.String
     type: TodoType
     created: dates.DateTime
-    is_deleted: bool
-    picture: Optional[bytes] = None
-    is_favorite: bool = False
+    is_deleted: _typedefs.Bool
+    picture: Optional[_typedefs.Binary] = None
+    is_favorite: _typedefs.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)
+    todos: Dict[_typedefs.I32, TodoItem] = field(default_factory=dict)
+    plain_ids: Set[_typedefs.I32] = field(default_factory=lambda: {1, 2, 3})
+    note_ids: List[_typedefs.I32] = field(default_factory=list)
+    checkboxes_ids: Set[_typedefs.I32] = field(default_factory=set)
 
 default_created_date: dates.DateTime = dates.DateTime(
     year=2024, month=12, day=25, hour=0, minute=0, second=0, microsecond=0
 )
 
 class Todo:
-    def create(self, text: str, type: TodoType) -> int: ...
+    def create(self, text: _typedefs.String, type: TodoType) -> _typedefs.I32: ...
     def update(self, item: TodoItem) -> None: ...
-    def get(self, id: int) -> TodoItem: ...
+    def get(self, id: _typedefs.I32) -> 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 filter(self, ids: List[_typedefs.I32]) -> List[TodoItem]: ...
+    def stats(self) -> Dict[_typedefs.I32, _typedefs.Double]: ...
+    def types(self) -> Set[_typedefs.I16]: ...
     def groupby(self) -> Dict[TodoType, List[TodoItem]]: ...
-    def ping(self) -> str: ...
+    def ping(self) -> _typedefs.String: ...
diff --git a/example/app/interfaces/todo_v2.pyi b/example/app/interfaces/todo_v2.pyi
index 78a8002..16fc41c 100644
--- a/example/app/interfaces/todo_v2.pyi
+++ b/example/app/interfaces/todo_v2.pyi
@@ -1 +1,3 @@
+from . import _typedefs
+
 class Todo: ...
diff --git a/pyproject.toml b/pyproject.toml
index 592925b..c5d0bee 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "thrift-pyi"
-version = "1.2.0"
+version = "2.0.0"
 description = "This is simple `.pyi` stubs generator from thrift interfaces"
 readme = "README.rst"
 repository = "https://github.com/unmade/thrift-pyi"
diff --git a/src/thriftpyi/main.py b/src/thriftpyi/main.py
index 3b29bf3..2fea061 100644
--- a/src/thriftpyi/main.py
+++ b/src/thriftpyi/main.py
@@ -24,6 +24,9 @@ def thriftpyi(  # pylint: disable=too-many-locals
     stub = stubs.build_init(path.stem for path in interfaces)
     files.save(ast.unparse(stub), to=output_dir / "__init__.pyi")
 
+    stub = stubs.build_typedefs()
+    files.save(ast.unparse(stub), to=output_dir / "_typedefs.pyi")
+
     for path in interfaces:
         tmodule = thriftpy2.load(str(path))
         proxy = TModuleProxy(tmodule)
diff --git a/src/thriftpyi/stubs.py b/src/thriftpyi/stubs.py
index 2740909..91f5499 100644
--- a/src/thriftpyi/stubs.py
+++ b/src/thriftpyi/stubs.py
@@ -30,6 +30,83 @@ def build_init(imports: Iterable[str]) -> ast.Module:
     )
 
 
+def build_typedefs() -> ast.Module:
+    types_map = {
+        "Binary": ("bytes", "BINARY"),
+        "Bool": ("bool", "BOOL"),
+        "Byte": ("int", "BYTE"),
+        "Double": ("float", "DOUBLE"),
+        "I16": ("int", "I16"),
+        "I32": ("int", "I32"),
+        "I64": ("int", "I64"),
+        "String": ("str", "STRING"),
+    }
+
+    body: list[ast.stmt] = [
+        ast.ImportFrom(
+            module="dataclasses",
+            names=[ast.alias(name="dataclass")],
+            level=0,
+        ),
+        ast.ImportFrom(
+            module="typing",
+            names=[ast.alias(name="Annotated")],
+            level=0,
+        ),
+        ast.ClassDef(
+            name="Metadata",
+            bases=[],
+            keywords=[],
+            body=[
+                ast.AnnAssign(
+                    target=ast.Name(id="thrift_type"),
+                    annotation=ast.Name(id="str"),
+                    simple=1,
+                ),
+            ],
+            decorator_list=[
+                ast.Call(
+                    func=ast.Name(id="dataclass"),
+                    args=[],
+                    keywords=[
+                        ast.keyword(
+                            arg="frozen",
+                            value=ast.Constant(value=True),
+                        ),
+                    ],
+                ),
+            ],
+        ),
+    ]
+
+    for alias, (base_type, thrift_type) in types_map.items():
+        assign = ast.Assign(
+            targets=[ast.Name(id=alias)],
+            value=ast.Subscript(
+                value=ast.Name(id="Annotated"),
+                slice=ast.Tuple(
+                    elts=[
+                        ast.Name(id=base_type),
+                        ast.Call(
+                            func=ast.Name(id="Metadata"),
+                            args=[],
+                            keywords=[
+                                ast.keyword(
+                                    arg="thrift_type",
+                                    value=ast.Constant(value=thrift_type),
+                                ),
+                            ],
+                        ),
+                    ],
+                ),
+            ),
+            lineno=0,
+        )
+        body.append(assign)
+
+    return ast.Module(body=body, type_ignores=[])
+
+
 def _make_imports(proxy: TModuleProxy) -> list[ast.ImportFrom]:
     imports = []
     if proxy.has_structs():
@@ -38,6 +115,7 @@ def _make_imports(proxy: TModuleProxy) -> list[ast.ImportFrom]:
         imports.append(_make_absolute_import("enum", "IntEnum"))
 
     imports.append(_make_absolute_import("typing", "*"))
+    imports.append(_make_relative_import(["_typedefs"]))
     imports.extend(
         _make_relative_import([name]) for name in sorted(proxy.get_imports())
     )
diff --git a/src/thriftpyi/utils.py b/src/thriftpyi/utils.py
index aa3a823..cda2815 100644
--- a/src/thriftpyi/utils.py
+++ b/src/thriftpyi/utils.py
@@ -47,43 +47,43 @@ def get_python_type(ttype: int, meta: List) -> str:
 
 def _get_bool(meta: List) -> str:
     del meta
-    return "bool"
+    return "_typedefs.Bool"
 
 
 def _get_double(meta: List) -> str:
     del meta
-    return "float"
+    return "_typedefs.Double"
 
 
 def _get_byte(meta: List) -> str:
     del meta
-    return "int"
+    return "_typedefs.Byte"
 
 
 def _get_binary(meta: List) -> str:
     del meta
-    return "bytes"
+    return "_typedefs.Binary"
 
 
 def _get_i16(meta: List) -> str:
     del meta
-    return "int"
+    return "_typedefs.I16"
 
 
 def _get_i32(meta: List) -> str:
     if meta and meta[0] is not None:
         return f"{meta[0].__module__}.{meta[0].__name__}"
-    return "int"
+    return "_typedefs.I32"
 
 
 def _get_i64(meta: List) -> str:
     del meta
-    return "int"
+    return "_typedefs.I64"
 
 
 def _get_str(meta: List) -> str:
     del meta
-    return "str"
+    return "_typedefs.String"
 
 
 def _get_struct(meta: List) -> str:


Test patch:
diff --git a/tests/stubs/expected/async/_typedefs.pyi b/tests/stubs/expected/async/_typedefs.pyi
new file mode 100644
index 0000000..823ef75
--- /dev/null
+++ b/tests/stubs/expected/async/_typedefs.pyi
@@ -0,0 +1,15 @@
+from dataclasses import dataclass
+from typing import Annotated
+
+@dataclass(frozen=True)
+class Metadata:
+    thrift_type: str
+
+Binary = Annotated[bytes, Metadata(thrift_type="BINARY")]
+Bool = Annotated[bool, Metadata(thrift_type="BOOL")]
+Byte = Annotated[int, Metadata(thrift_type="BYTE")]
+Double = Annotated[float, Metadata(thrift_type="DOUBLE")]
+I16 = Annotated[int, Metadata(thrift_type="I16")]
+I32 = Annotated[int, Metadata(thrift_type="I32")]
+I64 = Annotated[int, Metadata(thrift_type="I64")]
+String = Annotated[str, Metadata(thrift_type="STRING")]
diff --git a/tests/stubs/expected/async/dates.pyi b/tests/stubs/expected/async/dates.pyi
index 83e18c2..c0a1f51 100644
--- a/tests/stubs/expected/async/dates.pyi
+++ b/tests/stubs/expected/async/dates.pyi
@@ -1,15 +1,16 @@
 from dataclasses import dataclass
 from typing import *
+from . import _typedefs
 
 @dataclass
 class DateTime:
-    year: int
-    month: int
-    day: int
-    hour: int
-    minute: int
-    second: int
-    microsecond: Optional[int] = None
+    year: _typedefs.I16
+    month: _typedefs.Byte
+    day: _typedefs.Byte
+    hour: _typedefs.I16
+    minute: _typedefs.Byte
+    second: _typedefs.Byte
+    microsecond: Optional[_typedefs.I64] = None
 
 @dataclass
 class Date: ...
diff --git a/tests/stubs/expected/async/shared.pyi b/tests/stubs/expected/async/shared.pyi
index 90cae81..143cace 100644
--- a/tests/stubs/expected/async/shared.pyi
+++ b/tests/stubs/expected/async/shared.pyi
@@ -1,21 +1,22 @@
 from dataclasses import dataclass
 from typing import *
+from . import _typedefs
 
 class NotFound(Exception):
-    message: Optional[str] = "Not Found"
+    message: Optional[_typedefs.String] = "Not Found"
 
-    def __init__(self, message: Optional[str] = "Not Found") -> None: ...
+    def __init__(self, message: Optional[_typedefs.String] = "Not Found") -> None: ...
 
 class EmptyException(Exception): ...
 
 @dataclass
 class LimitOffset:
-    limit: Optional[int] = None
-    offset: Optional[int] = None
+    limit: Optional[_typedefs.I32] = None
+    offset: Optional[_typedefs.I32] = 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: ...
+    async def ping(self) -> _typedefs.String: ...
diff --git a/tests/stubs/expected/async/todo.pyi b/tests/stubs/expected/async/todo.pyi
index dfe85a3..409b5ec 100644
--- a/tests/stubs/expected/async/todo.pyi
+++ b/tests/stubs/expected/async/todo.pyi
@@ -1,6 +1,7 @@
 from dataclasses import dataclass, field
 from enum import IntEnum
 from typing import *
+from . import _typedefs
 from . import dates
 from . import shared
 
@@ -11,32 +12,32 @@ class TodoType(IntEnum):
 
 @dataclass
 class TodoItem:
-    id: int
-    text: str
+    id: _typedefs.I32
+    text: _typedefs.String
     type: TodoType
     created: dates.DateTime
-    is_deleted: bool
-    picture: Optional[bytes] = None
-    is_favorite: bool = False
+    is_deleted: _typedefs.Bool
+    picture: Optional[_typedefs.Binary] = None
+    is_favorite: _typedefs.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)
+    todos: Dict[_typedefs.I32, TodoItem] = field(default_factory=dict)
+    plain_ids: Set[_typedefs.I32] = field(default_factory=lambda: {1, 2, 3})
+    note_ids: List[_typedefs.I32] = field(default_factory=list)
+    checkboxes_ids: Set[_typedefs.I32] = field(default_factory=set)
 
 default_created_date: dates.DateTime = dates.DateTime(
     year=2024, month=12, day=25, hour=0, minute=0, second=0, microsecond=0
 )
 
 class Todo:
-    async def create(self, text: str, type: TodoType) -> int: ...
+    async def create(self, text: _typedefs.String, type: TodoType) -> _typedefs.I32: ...
     async def update(self, item: TodoItem) -> None: ...
-    async def get(self, id: int) -> TodoItem: ...
+    async def get(self, id: _typedefs.I32) -> TodoItem: ...
     async def all(self, pager: shared.LimitOffset) -> List[TodoItem]: ...
-    async def filter(self, ids: List[int]) -> List[TodoItem]: ...
-    async def stats(self) -> Dict[int, float]: ...
-    async def types(self) -> Set[int]: ...
+    async def filter(self, ids: List[_typedefs.I32]) -> List[TodoItem]: ...
+    async def stats(self) -> Dict[_typedefs.I32, _typedefs.Double]: ...
+    async def types(self) -> Set[_typedefs.I16]: ...
     async def groupby(self) -> Dict[TodoType, List[TodoItem]]: ...
-    async def ping(self) -> str: ...
+    async def ping(self) -> _typedefs.String: ...
diff --git a/tests/stubs/expected/async/todo_v2.pyi b/tests/stubs/expected/async/todo_v2.pyi
index 78a8002..16fc41c 100644
--- a/tests/stubs/expected/async/todo_v2.pyi
+++ b/tests/stubs/expected/async/todo_v2.pyi
@@ -1 +1,3 @@
+from . import _typedefs
+
 class Todo: ...
diff --git a/tests/stubs/expected/optional/_typedefs.pyi b/tests/stubs/expected/optional/_typedefs.pyi
new file mode 100644
index 0000000..823ef75
--- /dev/null
+++ b/tests/stubs/expected/optional/_typedefs.pyi
@@ -0,0 +1,15 @@
+from dataclasses import dataclass
+from typing import Annotated
+
+@dataclass(frozen=True)
+class Metadata:
+    thrift_type: str
+
+Binary = Annotated[bytes, Metadata(thrift_type="BINARY")]
+Bool = Annotated[bool, Metadata(thrift_type="BOOL")]
+Byte = Annotated[int, Metadata(thrift_type="BYTE")]
+Double = Annotated[float, Metadata(thrift_type="DOUBLE")]
+I16 = Annotated[int, Metadata(thrift_type="I16")]
+I32 = Annotated[int, Metadata(thrift_type="I32")]
+I64 = Annotated[int, Metadata(thrift_type="I64")]
+String = Annotated[str, Metadata(thrift_type="STRING")]
diff --git a/tests/stubs/expected/optional/dates.pyi b/tests/stubs/expected/optional/dates.pyi
index 7f4bc78..02be390 100644
--- a/tests/stubs/expected/optional/dates.pyi
+++ b/tests/stubs/expected/optional/dates.pyi
@@ -1,15 +1,16 @@
 from dataclasses import dataclass
 from typing import *
+from . import _typedefs
 
 @dataclass
 class DateTime:
-    year: Optional[int] = None
-    month: Optional[int] = None
-    day: Optional[int] = None
-    hour: Optional[int] = None
-    minute: Optional[int] = None
-    second: Optional[int] = None
-    microsecond: Optional[int] = None
+    year: Optional[_typedefs.I16] = None
+    month: Optional[_typedefs.Byte] = None
+    day: Optional[_typedefs.Byte] = None
+    hour: Optional[_typedefs.I16] = None
+    minute: Optional[_typedefs.Byte] = None
+    second: Optional[_typedefs.Byte] = None
+    microsecond: Optional[_typedefs.I64] = None
 
 @dataclass
 class Date: ...
diff --git a/tests/stubs/expected/optional/shared.pyi b/tests/stubs/expected/optional/shared.pyi
index ff9725a..f6bbd56 100644
--- a/tests/stubs/expected/optional/shared.pyi
+++ b/tests/stubs/expected/optional/shared.pyi
@@ -1,21 +1,22 @@
 from dataclasses import dataclass
 from typing import *
+from . import _typedefs
 
 class NotFound(Exception):
-    message: Optional[str] = "Not Found"
+    message: Optional[_typedefs.String] = "Not Found"
 
-    def __init__(self, message: Optional[str] = "Not Found") -> None: ...
+    def __init__(self, message: Optional[_typedefs.String] = "Not Found") -> None: ...
 
 class EmptyException(Exception): ...
 
 @dataclass
 class LimitOffset:
-    limit: Optional[int] = None
-    offset: Optional[int] = None
+    limit: Optional[_typedefs.I32] = None
+    offset: Optional[_typedefs.I32] = 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: ...
+    def ping(self) -> _typedefs.String: ...
diff --git a/tests/stubs/expected/optional/todo.pyi b/tests/stubs/expected/optional/todo.pyi
index 5798956..baa5a50 100644
--- a/tests/stubs/expected/optional/todo.pyi
+++ b/tests/stubs/expected/optional/todo.pyi
@@ -1,6 +1,7 @@
 from dataclasses import dataclass, field
 from enum import IntEnum
 from typing import *
+from . import _typedefs
 from . import dates
 from . import shared
 
@@ -11,32 +12,32 @@ class TodoType(IntEnum):
 
 @dataclass
 class TodoItem:
-    id: Optional[int] = None
-    text: Optional[str] = None
+    id: Optional[_typedefs.I32] = None
+    text: Optional[_typedefs.String] = None
     type: Optional[TodoType] = None
     created: Optional[dates.DateTime] = None
-    is_deleted: Optional[bool] = None
-    picture: Optional[bytes] = None
-    is_favorite: Optional[bool] = False
+    is_deleted: Optional[_typedefs.Bool] = None
+    picture: Optional[_typedefs.Binary] = None
+    is_favorite: Optional[_typedefs.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)
+    todos: Optional[Dict[_typedefs.I32, TodoItem]] = field(default_factory=dict)
+    plain_ids: Optional[Set[_typedefs.I32]] = field(default_factory=lambda: {1, 2, 3})
+    note_ids: Optional[List[_typedefs.I32]] = field(default_factory=list)
+    checkboxes_ids: Optional[Set[_typedefs.I32]] = field(default_factory=set)
 
 default_created_date: dates.DateTime = dates.DateTime(
     year=2024, month=12, day=25, hour=0, minute=0, second=0, microsecond=0
 )
 
 class Todo:
-    def create(self, text: str, type: TodoType) -> int: ...
+    def create(self, text: _typedefs.String, type: TodoType) -> _typedefs.I32: ...
     def update(self, item: TodoItem) -> None: ...
-    def get(self, id: int) -> TodoItem: ...
+    def get(self, id: _typedefs.I32) -> 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 filter(self, ids: List[_typedefs.I32]) -> List[TodoItem]: ...
+    def stats(self) -> Dict[_typedefs.I32, _typedefs.Double]: ...
+    def types(self) -> Set[_typedefs.I16]: ...
     def groupby(self) -> Dict[TodoType, List[TodoItem]]: ...
-    def ping(self) -> str: ...
+    def ping(self) -> _typedefs.String: ...
diff --git a/tests/stubs/expected/optional/todo_v2.pyi b/tests/stubs/expected/optional/todo_v2.pyi
index 78a8002..16fc41c 100644
--- a/tests/stubs/expected/optional/todo_v2.pyi
+++ b/tests/stubs/expected/optional/todo_v2.pyi
@@ -1 +1,3 @@
+from . import _typedefs
+
 class Todo: ...
diff --git a/tests/stubs/expected/sync/_typedefs.pyi b/tests/stubs/expected/sync/_typedefs.pyi
new file mode 100644
index 0000000..823ef75
--- /dev/null
+++ b/tests/stubs/expected/sync/_typedefs.pyi
@@ -0,0 +1,15 @@
+from dataclasses import dataclass
+from typing import Annotated
+
+@dataclass(frozen=True)
+class Metadata:
+    thrift_type: str
+
+Binary = Annotated[bytes, Metadata(thrift_type="BINARY")]
+Bool = Annotated[bool, Metadata(thrift_type="BOOL")]
+Byte = Annotated[int, Metadata(thrift_type="BYTE")]
+Double = Annotated[float, Metadata(thrift_type="DOUBLE")]
+I16 = Annotated[int, Metadata(thrift_type="I16")]
+I32 = Annotated[int, Metadata(thrift_type="I32")]
+I64 = Annotated[int, Metadata(thrift_type="I64")]
+String = Annotated[str, Metadata(thrift_type="STRING")]
diff --git a/tests/stubs/expected/sync/dates.pyi b/tests/stubs/expected/sync/dates.pyi
index 83e18c2..c0a1f51 100644
--- a/tests/stubs/expected/sync/dates.pyi
+++ b/tests/stubs/expected/sync/dates.pyi
@@ -1,15 +1,16 @@
 from dataclasses import dataclass
 from typing import *
+from . import _typedefs
 
 @dataclass
 class DateTime:
-    year: int
-    month: int
-    day: int
-    hour: int
-    minute: int
-    second: int
-    microsecond: Optional[int] = None
+    year: _typedefs.I16
+    month: _typedefs.Byte
+    day: _typedefs.Byte
+    hour: _typedefs.I16
+    minute: _typedefs.Byte
+    second: _typedefs.Byte
+    microsecond: Optional[_typedefs.I64] = None
 
 @dataclass
 class Date: ...
diff --git a/tests/stubs/expected/sync/shared.pyi b/tests/stubs/expected/sync/shared.pyi
index ff9725a..f6bbd56 100644
--- a/tests/stubs/expected/sync/shared.pyi
+++ b/tests/stubs/expected/sync/shared.pyi
@@ -1,21 +1,22 @@
 from dataclasses import dataclass
 from typing import *
+from . import _typedefs
 
 class NotFound(Exception):
-    message: Optional[str] = "Not Found"
+    message: Optional[_typedefs.String] = "Not Found"
 
-    def __init__(self, message: Optional[str] = "Not Found") -> None: ...
+    def __init__(self, message: Optional[_typedefs.String] = "Not Found") -> None: ...
 
 class EmptyException(Exception): ...
 
 @dataclass
 class LimitOffset:
-    limit: Optional[int] = None
-    offset: Optional[int] = None
+    limit: Optional[_typedefs.I32] = None
+    offset: Optional[_typedefs.I32] = 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: ...
+    def ping(self) -> _typedefs.String: ...
diff --git a/tests/stubs/expected/sync/todo.pyi b/tests/stubs/expected/sync/todo.pyi
index 27bc446..f520042 100644
--- a/tests/stubs/expected/sync/todo.pyi
+++ b/tests/stubs/expected/sync/todo.pyi
@@ -1,6 +1,7 @@
 from dataclasses import dataclass, field
 from enum import IntEnum
 from typing import *
+from . import _typedefs
 from . import dates
 from . import shared
 
@@ -11,32 +12,32 @@ class TodoType(IntEnum):
 
 @dataclass
 class TodoItem:
-    id: int
-    text: str
+    id: _typedefs.I32
+    text: _typedefs.String
     type: TodoType
     created: dates.DateTime
-    is_deleted: bool
-    picture: Optional[bytes] = None
-    is_favorite: bool = False
+    is_deleted: _typedefs.Bool
+    picture: Optional[_typedefs.Binary] = None
+    is_favorite: _typedefs.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)
+    todos: Dict[_typedefs.I32, TodoItem] = field(default_factory=dict)
+    plain_ids: Set[_typedefs.I32] = field(default_factory=lambda: {1, 2, 3})
+    note_ids: List[_typedefs.I32] = field(default_factory=list)
+    checkboxes_ids: Set[_typedefs.I32] = field(default_factory=set)
 
 default_created_date: dates.DateTime = dates.DateTime(
     year=2024, month=12, day=25, hour=0, minute=0, second=0, microsecond=0
 )
 
 class Todo:
-    def create(self, text: str, type: TodoType) -> int: ...
+    def create(self, text: _typedefs.String, type: TodoType) -> _typedefs.I32: ...
     def update(self, item: TodoItem) -> None: ...
-    def get(self, id: int) -> TodoItem: ...
+    def get(self, id: _typedefs.I32) -> 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 filter(self, ids: List[_typedefs.I32]) -> List[TodoItem]: ...
+    def stats(self) -> Dict[_typedefs.I32, _typedefs.Double]: ...
+    def types(self) -> Set[_typedefs.I16]: ...
     def groupby(self) -> Dict[TodoType, List[TodoItem]]: ...
-    def ping(self) -> str: ...
+    def ping(self) -> _typedefs.String: ...
diff --git a/tests/stubs/expected/sync/todo_v2.pyi b/tests/stubs/expected/sync/todo_v2.pyi
index 78a8002..16fc41c 100644
--- a/tests/stubs/expected/sync/todo_v2.pyi
+++ b/tests/stubs/expected/sync/todo_v2.pyi
@@ -1 +1,3 @@
+from . import _typedefs
+
 class Todo: ...
diff --git a/tests/test_cli.py b/tests/test_cli.py
index a31bd96..3db11ce 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -22,7 +22,14 @@ def test_main(capsys, expected_dir, args):
 
     main([input_dir, "--output", output_dir, *args])
 
-    pyi_files = ["__init__.pyi", "dates.pyi", "shared.pyi", "todo.pyi", "todo_v2.pyi"]
+    pyi_files = [
+        "__init__.pyi",
+        "_typedefs.pyi",
+        "dates.pyi",
+        "shared.pyi",
+        "todo.pyi",
+        "todo_v2.pyi",
+    ]
     match, mismatch, errors = filecmp.cmpfiles(output_dir, expected_dir, pyi_files)
     assert errors == []
     assert mismatch == []