Skip to content

Commit

Permalink
Dump consts after structs (#47)
Browse files Browse the repository at this point in the history
  • Loading branch information
unmade committed Apr 27, 2024
1 parent 8fde57c commit dc4d139
Show file tree
Hide file tree
Showing 10 changed files with 64 additions and 18 deletions.
2 changes: 1 addition & 1 deletion .pylintrc
Original file line number Diff line number Diff line change
Expand Up @@ -459,7 +459,7 @@ max-parents=7
max-public-methods=10

# Maximum number of return / yield for function / method body.
max-returns=4
max-returns=6

# Maximum number of statements in function / method body.
max-statements=20
Expand Down
12 changes: 11 additions & 1 deletion example/interfaces/dates.thrift
Original file line number Diff line number Diff line change
Expand Up @@ -12,4 +12,14 @@ struct DateTime {

struct Date {

}
}

const DateTime EPOCH = {
"year": 1970,
"month": 1,
"day": 1,
"hour": 0,
"minute": 0,
"second": 0,
"microsecond": 0
}
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[tool.poetry]
name = "thrift-pyi"
version = "0.8.0"
version = "0.9.0"
description = "This is simple `.pyi` stubs generator from thrift interfaces"
readme = "README.rst"
repository = "https://github.com/unmade/thrift-pyi"
Expand Down
12 changes: 9 additions & 3 deletions src/thriftpyi/proxies.py
Original file line number Diff line number Diff line change
Expand Up @@ -52,12 +52,18 @@ def has_structs(self) -> bool:
def has_enums(self) -> bool:
return len(self.tmodule.__thrift_meta__["enums"]) > 0

@staticmethod
def _make_const(tconst) -> Field:
def _make_const(self, tconst) -> Field:
name, value = tconst
return Field(
name=name,
type=guess_type(value),
type=guess_type(
value,
known_modules={
module.__name__
for module in self.tmodule.__thrift_meta__["includes"]
},
known_structs=self.tmodule.__thrift_meta__["structs"],
),
value=value,
required=True,
)
Expand Down
2 changes: 1 addition & 1 deletion src/thriftpyi/stubs.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,10 @@ def build(
return ast.Module(
body=[
*_make_imports(proxy),
*_make_consts(proxy),
*_make_exceptions(proxy, strict=strict_fields),
*_make_enums(proxy),
*_make_structs(proxy, strict=strict_fields),
*_make_consts(proxy),
*_make_service(proxy, is_async, strict=strict_methods),
],
type_ignores=[],
Expand Down
34 changes: 28 additions & 6 deletions src/thriftpyi/utils.py
Original file line number Diff line number Diff line change
@@ -1,21 +1,43 @@
from collections.abc import Collection, Mapping
from typing import List, Tuple
from collections.abc import Mapping
from typing import Any, Collection, List, Tuple, Type

from thriftpy2.thrift import TType


def guess_type(value) -> str:
def guess_type( # pylint: disable=too-many-branches
value, *, known_modules: Collection[str], known_structs: Collection[Type[Any]]
) -> str:
if isinstance(value, (bool, int, float, str, bytes)):
return type(value).__name__

if isinstance(value, Mapping):
type_ = type(value).__name__.capitalize()
key_type = guess_type(list(value.keys())[0])
value_type = guess_type(list(value.values())[0])
key_type = guess_type(
next(iter(value.keys())),
known_modules=known_modules,
known_structs=known_structs,
)
value_type = guess_type(
next(iter(value.values())),
known_modules=known_modules,
known_structs=known_structs,
)
return f"{type_}[{key_type}, {value_type}]"

if isinstance(value, Collection):
type_ = type(value).__name__.capitalize()
item_type = guess_type(next(iter(value)))
item_type = guess_type(
next(iter(value)), known_modules=known_modules, known_structs=known_structs
)
return f"{type_}[{item_type}]"

if hasattr(value, "__class__"):
module_name: str = value.__class__.__module__
class_name: str = value.__class__.__name__
if module_name in known_modules:
return f"{module_name}.{class_name}"
if type(value) in known_structs:
return class_name
return "Any"


Expand Down
4 changes: 4 additions & 0 deletions tests/stubs/expected/optional/dates.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
4 changes: 4 additions & 0 deletions tests/stubs/expected/sync/dates.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
8 changes: 4 additions & 4 deletions tests/stubs/expected/sync/shared.pyi
Original file line number Diff line number Diff line change
@@ -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"

Expand All @@ -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: ...
2 changes: 1 addition & 1 deletion tests/test_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ class TestGuessType:
],
)
def test(self, guess_type, value, expected: str):
assert guess_type(value) == expected
assert guess_type(value, known_modules=[], known_structs=[]) == expected


class TestRegisterBinary:
Expand Down

0 comments on commit dc4d139

Please sign in to comment.