Skip to content

Commit

Permalink
Minor refactoring of mypy_plugin
Browse files Browse the repository at this point in the history
  • Loading branch information
sg495 committed Feb 23, 2024
1 parent f082935 commit dcfb605
Showing 1 changed file with 9 additions and 12 deletions.
21 changes: 9 additions & 12 deletions typed_descriptors/mypy_plugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,33 +6,31 @@
from __future__ import annotations
from collections.abc import Callable
import typing
from mypy.types import CallableType, Instance, get_proper_type, Type
from mypy.plugin import FunctionContext, Plugin
from mypy.types import CallableType, Instance, get_proper_type, Type


def typed_descriptor_ret_type_arg0(ctx: FunctionContext) -> Type:
def typed_descriptor_hook(ctx: FunctionContext) -> Type:
"""
Extracts the descriptor type and assigns it to the generic type parameter,
presuming that the generic type parameter appears at index 0.
"""
assert ctx.arg_types and ctx.arg_types[0]
descriptor_t_constr = get_proper_type(ctx.arg_types[0][0])
_descriptor_t = get_proper_type(ctx.arg_types[0][0])
ret_t = get_proper_type(ctx.default_return_type)
if isinstance(descriptor_t_constr, CallableType):
descriptor_t = descriptor_t_constr.ret_type
if isinstance(_descriptor_t, CallableType):
descriptor_t = _descriptor_t.ret_type
if isinstance(ret_t, Instance):
args = list(ret_t.args)
args[0] = descriptor_t
ret_t = ret_t.copy_modified(args=tuple(args))
return ret_t.copy_modified(args=tuple(args))
return ret_t
return ret_t


_function_hooks = {
"typed_descriptors.attr.Attr": typed_descriptor_ret_type_arg0,
"typed_descriptors.prop.Prop": typed_descriptor_ret_type_arg0,
"typed_descriptors.attr.Attr": typed_descriptor_hook,
"typed_descriptors.prop.Prop": typed_descriptor_hook,
}


class TypedDescriptorsPlugin(Plugin):
"""
Mypy plugin which expands type inference for typed descriptor,
Expand All @@ -47,7 +45,6 @@ def get_function_hook(
return hook
return super().get_function_hook(fullname)


def plugin(version: str) -> typing.Type[TypedDescriptorsPlugin]:
"""
Entry point for the Mypy plugin.
Expand Down

0 comments on commit dcfb605

Please sign in to comment.