def def_bool_method(name: str) -> None:
_define_method(context, enum_info, context.type.name, name, [
nodes.Argument(nodes.Var('self', enum_instance), enum_instance,
None, nodes.ARG_POS),
nodes.Argument(nodes.Var('other', enum_instance), enum_instance,
None, nodes.ARG_POS),
], bool_type)
def newmodel_hook(ctx: ClassDefContext) -> None:
"""Add a literal type for _name.
This makes tagged unions possible. If we have Union[ResUsers, ResPartner]
then mypy understands that `if x._name == "res.users":` means that x is
ResUsers.
Unfortunately it doesn't work for narrowing down BaseModel, so it's only
rarely useful. A workaround could be to automatically define a union of all
defined models.
PEP 622 currently proposes @typing.sealed, which would mandate that all
subclasses of a class must be mandated in the same module. This should
resolve the problem for the current type stub setup, and may help with a
version of the plugin that can analyze actual Odoo code.
https://www.python.org/dev/peps/pep-0622/#sealed-classes-as-adts
"""
dotted_name = [] # type: t.List[str]
for char in ctx.cls.name:
if dotted_name and char.isupper():
dotted_name.append(".")
dotted_name.append(char.lower())
name_type = types.LiteralType("".join(dotted_name),
ctx.api.named_type("str"))
var = nodes.Var("_name", name_type)
var.info = ctx.cls.info
stn = nodes.SymbolTableNode(nodes.MDEF, var)
ctx.cls.info.names["_name"] = stn
def _transform(self) -> List[Field]:
ctx = self._ctx
fields = self._collect_fields()
schema_t = self._lookup_type('edb.schema.schema.Schema')
for f in fields:
if f.has_explicit_accessor:
continue
mypy_helpers.add_method(
ctx,
name=f'get_{f.name}',
args=[
nodes.Argument(
variable=nodes.Var(
name='schema',
type=schema_t,
),
type_annotation=schema_t,
initializer=None,
kind=nodes.ARG_POS,
),
],
return_type=f.type,
)
return fields
def add_model_replace(ctx) -> None:
"""Add model replace method."""
any_type = types.AnyType(types.TypeOfAny.special_form)
var = nodes.Var('change', any_type)
kw_arg = nodes.Argument(variable=var,
type_annotation=any_type,
initializer=None,
kind=nodes.ARG_STAR2)
ret_type = types.Instance(ctx.cls.info, [])
common.add_method(ctx, 'replace', [kw_arg], ret_type)
def add_model_astuple(ctx) -> None:
"""Add model astuple method."""
bool_type = ctx.api.builtin_type('builtins.bool')
tuple_type = ctx.api.builtin_type('builtins.tuple')
var = nodes.Var('recurse', bool_type)
recurse = nodes.Argument(variable=var,
type_annotation=bool_type,
initializer=nodes.NameExpr('True'),
kind=nodes.ARG_NAMED_OPT)
common.add_method(ctx, 'astuple', [recurse], tuple_type)
def transform(self):
ctx = self._ctx
metadata = ctx.cls.info.metadata.get(METADATA_KEY)
if not metadata:
ctx.cls.info.metadata[METADATA_KEY] = metadata = {}
metadata['processing'] = True
if metadata.get('processed'):
return
try:
fields = self._collect_fields()
schema_t = self._lookup_type('edb.schema.schema.Schema')
except DeferException:
ctx.api.defer()
return None
cls_info = ctx.cls.info
for f in fields:
ftype = cls_info.get(f.name).type
if ftype is None or cls_info.get(f'get_{f.name}') is not None:
# The class is already doing something funny with the
# field or the accessor, so ignore it.
continue
if f.is_optional:
ftype = types.UnionType.make_union(
[ftype, types.NoneType()],
line=ftype.line,
column=ftype.column,
)
mypy_helpers.add_method(
ctx,
name=f'get_{f.name}',
args=[
nodes.Argument(
variable=nodes.Var(
name='schema',
type=schema_t,
),
type_annotation=schema_t,
initializer=None,
kind=nodes.ARG_POS,
),
],
return_type=ftype,
)
metadata['fields'] = {f.name: f.serialize() for f in fields}
metadata['processed'] = True
def add_model_set(ctx) -> None:
"""Add model fields method."""
args = []
str_type = ctx.api.builtin_type('builtins.str')
name_var = nodes.Var('name', str_type)
name_arg = nodes.Argument(variable=name_var,
type_annotation=str_type,
initializer=None,
kind=nodes.ARG_POS)
args.append(name_arg)
any_type = types.AnyType(types.TypeOfAny.special_form)
value_var = nodes.Var('value', any_type)
value_arg = nodes.Argument(variable=value_var,
type_annotation=any_type,
initializer=None,
kind=nodes.ARG_POS)
args.append(value_arg)
common.add_method(ctx, '_set', args, types.NoneTyp())
def add_model_init(ctx, var_types, var_fields) -> None:
"""Add dummy init method to class."""
args = []
for (var_type, var_field) in zip(var_types, var_fields):
var = nodes.Var(var_field, var_type)
args.append(
nodes.Argument(variable=var,
type_annotation=var_type,
initializer=None,
kind=nodes.ARG_POS))
common.add_method(ctx, '__init__', args, types.NoneTyp())
def _add_member_to_class(self, member_name: str, member_type: types.Type,
clazz: nodes.TypeInfo) -> None:
"""Add a new member to the class.
Add a variable with given name and type to the symbol table of a
class. This also takes care about setting necessary attributes on the
variable node.
"""
var = nodes.Var(member_name)
var.info = clazz
var._fullname = clazz.fullname + "." + member_name
var.type = member_type
clazz.names[member_name] = nodes.SymbolTableNode(nodes.MDEF, var)
self.log("Defined o.vo field: %s.%s as %s" %
(clazz.fullname, member_name, member_type))
def transform(self):
ctx = self._ctx
metadata_key = self._get_metadata_key()
metadata = ctx.cls.info.metadata.get(metadata_key)
if not metadata:
ctx.cls.info.metadata[metadata_key] = metadata = {}
metadata['processing'] = True
if metadata.get('processed'):
return
try:
fields = self._collect_fields()
schema_t = self._lookup_type('edb.schema.schema.Schema')
except DeferException:
ctx.api.defer()
return None
for f in fields:
if f.has_explicit_accessor:
continue
mypy_helpers.add_method(
ctx,
name=f'get_{f.name}',
args=[
nodes.Argument(
variable=nodes.Var(
name='schema',
type=schema_t,
),
type_annotation=schema_t,
initializer=None,
kind=nodes.ARG_POS,
),
],
return_type=f.type,
)
metadata['fields'] = {f.name: f.serialize() for f in fields}
metadata['processed'] = True
def transform_enum_type(context: mypy.plugin.AnalyzeTypeContext) -> types.Type:
'''
This is registered as a handler of the :code:`get_type_analyze_hook` hook
See https://mypy.readthedocs.io/en/latest/extending_mypy.html#current-list-of-plugin-hooks for
more information about hooks.
This will be the first hook called in this plugin.
It allows us to change or alter type definitions as mypy sees them.
This is needed because our Enum uses :code:`EnumMeta` class which defines :code:`__new__`
and it shadows everything for mypy. So without this callback our type
visible by mypy would look like following in :code:`transform_enum_class_def`
hook handler. (Please read comment in :code:`transform_enum_class_def`).
::
ClassDef:3(
Color
FallbackToAny
AssignmentStmt:4(
NameExpr(RED [m])
IntExpr(1)
builtins.int)
AssignmentStmt:5(
NameExpr(GREEN [m])
IntExpr(2)
builtins.int)
AssignmentStmt:6(
NameExpr(BLUE [m])
IntExpr(3)
builtins.int))
TypeInfo(
Name(main.Color)
Bases(builtins.object)
Mro(main.Color, builtins.object)
Names(
BLUE (builtins.int)
GREEN (builtins.int)
RED (builtins.int)))
Please check docs to :code:`transform_enum_class_def` method
where you can se what these types look like when this hook
is used.
For example you can see that :code:`TypeInfo` is completely missing
- inheritance to Enum type
- metaclass definition
So it would be possible to change attribute types without this hook
but we have to add a metaclass for :code:`__iter__` and :code:`__getitem__` definitions.
To do this we have to describe what :code:`Enum` and :code:`EnumMeta` look like.
We did not find how to tell mypy to give us these definitions
so we described them manually in this part, and then properly
added inheritance to any :code:`Enum` inherited class to our fake :code:`Enum`.
'''
# Find references to some builtins which are often used
type_type = context.api.named_type('builtins.type')
object_type = context.api.named_type('builtins.object')
str_type = context.api.named_type('builtins.str')
bool_type = context.api.named_type('builtins.bool')
# Define meta class in fake module :code:`_fastenum`
# This is roughly equivalent to:
#
# class EnumMeta(builtins.type): pass
#
# in module :code:`_fastenum`.
# We have to define two things: :code:`ClassDef` an AST node and it's type definition using :code:`TypeInfo`
# :code:`ClassDef` defines only class syntactically all attributes and such will be defined in :code:`TypeInfo`
meta_cls = nodes.ClassDef('EnumMeta', nodes.Block([nodes.PassStmt()]), [],
[type_type])
meta_cls.fullname = '_fastenum.EnumMeta'
meta_info = nodes.TypeInfo(nodes.SymbolTable(), meta_cls, '_fastenum')
# We have to define inheritance again, mypy :code:`ClassDef` and :code:`TypeInfo`
# won't automatically share this information so we have to tell it again it is inherited from :code:`builtins.type`
meta_info.bases = [type_type]
# Last thing to get everything working is to define mro (method resolution order)
# without correctly specifying this mypy won't complain but it wont see any method or attributes
# defined in parents or even class itself.
# So we have to define class itself as :code:`meta_info` and **all of it's parents** (even indirect one)
# (this is not working in transitional fashion)
meta_info.mro = [meta_info, type_type.type, object_type.type]
# Define Enum class which is using EnumMeta as its metaclass in the fake :code:`_fastenum` module
# This is very similar to the previous definition and it is roughly equivalent to:
#
# class Enum(metaclass = EnumMeta): pass
#
# Notice that we still define :code:`builtins.object` as it's parent
# even if we don't have to do it in Python3, but we have to do it here!
enum_cls = nodes.ClassDef('Enum', nodes.Block([nodes.PassStmt()]), [],
[object_type], nodes.NameExpr('EnumMeta'))
enum_cls.fullname = '_fastenum.Enum'
enum_info = nodes.TypeInfo(nodes.SymbolTable(), enum_cls, '_fastenum')
# Same as before we have to define all parents (even :code:`builtins.object`)
enum_info.bases = [object_type]
enum_info.mro = [enum_info, object_type.type]
# New things in here are that we have to define the metaclass again in info
# I don't know why we have to define it on :code:`metaclass_type` and :code:`declared_metaclass`
# at the same time but mypy requires it that way, otherwise it ignores that metaclass
# and does not complain at all.
enum_info.metaclass_type = types.Instance(meta_info, [])
enum_info.declared_metaclass = types.Instance(meta_info, [])
# Add the attribute ``value`` to enum instances.
# Don't be scared by :code:`TypeOfAny`. mypy just has multiple types of Any.
# See :code:`TypeOfAny` definition (it is an enum with comments).
value_attribute = nodes.Var('value',
types.AnyType(types.TypeOfAny.explicit))
value_attribute.is_initialized_in_class = False
# As before we have to link our variable back to our class.
value_attribute.info = enum_info
enum_info.names['value'] = nodes.SymbolTableNode(nodes.MDEF,
value_attribute,
plugin_generated=True)
# Add the attribute ``name``` to enum instances.
value_attribute = nodes.Var('name', str_type)
value_attribute.is_initialized_in_class = False
value_attribute.info = enum_info
enum_info.names['name'] = nodes.SymbolTableNode(nodes.MDEF,
value_attribute,
plugin_generated=True)
#
# So after these few lines we end up with something like:
#
# module `_fastenum`:
#
# class EnumMeta(builtins.type): pass
#
# class Enum(metaclass = EnumMeta):
# name: str
# value: Any
#
# Prepare TypeVar, all these lines are just:
# _EnumMetaType = TypeVar('_EnumMetaType', bound = 'EnumMeta')
# We just have to describe expressions and definitions separately for mypy
meta_enum_instance = types.Instance(meta_info, [])
self_tvar_expr = nodes.TypeVarExpr(
'_EnumMetaType', f'{meta_info.fullname()}._EnumMetaType', [],
meta_enum_instance)
meta_info.names['_EnumMetaType'] = nodes.SymbolTableNode(
nodes.MDEF, self_tvar_expr)
self_tvar_def = types.TypeVarDef('_EnumMetaType',
f'{meta_info.fullname()}._EnumMetaType',
-1, [], meta_enum_instance)
self_tvar_type = types.TypeVarType(self_tvar_def)
# Define base __iter__ and __next__ for our meta class and use TypeVar `_EnumMetaType` as its return value
# so we can say its return value is bound to all children.
# See more comments about the definition in:
# - `transform_enum_class_def` handler
# - and `_define_method` docs + comments
_define_method(context, meta_info, context.type.name, '__iter__', [],
self_tvar_type)
_define_method(context, meta_info, context.type.name, '__next__', [],
self_tvar_type)
# Same way with __getitem__
_define_method(context, meta_info, context.type.name, '__getitem__', [
nodes.Argument(nodes.Var('cls', meta_enum_instance),
meta_enum_instance, None, nodes.ARG_POS),
nodes.Argument(nodes.Var('key', str_type), str_type, None,
nodes.ARG_POS),
], self_tvar_type)
# We also have to support constructor interface of enum, so when someone calls Enum('value').
# This is simply done by adding the `__init__` method with two arguments (self, value).
enum_instance = types.Instance(enum_info, [])
any_type = types.AnyType(types.TypeOfAny.explicit)
_define_method(context, enum_info, context.type.name, '__init__', [
nodes.Argument(nodes.Var('self', enum_instance), enum_instance, None,
nodes.ARG_POS),
nodes.Argument(nodes.Var('value', any_type), any_type, None,
nodes.ARG_POS),
], types.NoneTyp())
# Because enums can be used even in comparison expression like `A > B`
# we have to support these methods in our fake enum class too.
def def_bool_method(name: str) -> None:
_define_method(context, enum_info, context.type.name, name, [
nodes.Argument(nodes.Var('self', enum_instance), enum_instance,
None, nodes.ARG_POS),
nodes.Argument(nodes.Var('other', enum_instance), enum_instance,
None, nodes.ARG_POS),
], bool_type)
for name in ('le', 'eq', 'ne', 'ge', 'gt'):
def_bool_method(f'__{name}__')
#
# After all this we end up with:
#
# module `_fastenum`:
#
# class EnumMeta(builtins.type):
# _EnumMetaType = TypeVar('_EnumMetaType', bound = 'EnumMeta')
#
# def __iter__() -> _EnumMetaType: pass
# def __next__() -> _EnumMetaType: pass
# def __getitem__(cls: 'EnumMeta', key: str) -> 'EnumMeta': pass
#
#
# class Enum(metaclass = EnumMeta):
# name: str
# value: Any
#
# def __init__(self, value: Any) -> None: pass
#
# def __le__(self, other: Enum) -> bool: pass
# def __eq__(self, other: Enum) -> bool: pass
# def __ne__(self, other: Enum) -> bool: pass
# def __ge__(self, other: Enum) -> bool: pass
# def __gt__(self, other: Enum) -> bool: pass
#
# And we have to return new type for our `Enum` class which will be our new `Enum`
return types.Instance(enum_info, [])
def transform_enum_class_def(context: mypy.plugin.ClassDefContext) -> None:
'''
This is registered as a handler of the :code:`get_base_class_hook` hook
See https://mypy.readthedocs.io/en/latest/extending_mypy.html#current-list-of-plugin-hooks for
more information about hooks.
It gives us the ability to get the defined type and AST nodes representing how
mypy sees our class (inherited from Enum) and perform any changes we want.
This hook is called after :code:`get_type_analyze_hook` hook which is handled by :code:`transform_enum_type`
so the type of the super class is slightly altered after that.
But if we inspect :code:`context.cls` which is of type :code:`nodes.ClassDef` (an AST node)
we get something like:
::
ClassDef:3(
Color
BaseType(_fastenum.Enum)
AssignmentStmt:4(
NameExpr(RED [m])
IntExpr(1)
builtins.int)
AssignmentStmt:5(
NameExpr(GREEN [m])
IntExpr(2)
builtins.int)
AssignmentStmt:6(
NameExpr(BLUE [m])
IntExpr(3)
builtins.int))
And by inspecting :code:`context.cls.info` which is of type :code:`nodes.TypeInfo` - still an AST node
but one which defines the type of our class (where previous defined class definition) you'll get:
::
TypeInfo(
Name(main.Color)
Bases(_fastenum.Enum)
Mro(main.Color, _fastenum.Enum, builtins.object)
Names(
BLUE (builtins.int)
GREEN (builtins.int)
RED (builtins.int))
MetaclassType(_fastenum.EnumMeta))
From these you can see that mypy sees our class attributes as int, str or any value our enum has
but not as Enum class instances. So we have to update that.
'''
info = context.cls.info
# This is a hotfix for the built-in `enum.Enum` class
# which inherits from the default builtins.int
# and therefor our comparison methods are not compatible.
# To get over that we remove `int` base class from that class def.
info.bases = [
base for base in info.bases if base.type.fullname() != 'builtins.int'
]
# First clear all `nodes.AssignmentStmt` in class.
# These are basically class-level attributes defining enum values.
# This way we will remove the attribute assignment statement
# that mypy sees when the class is defined as:
#
# class Color(Enum):
# RED = 1
#
# From this mypy creates :code:`AssignmentStmt` where RED = IntInstance(1)
# We remove this, so mypy see only class variables in :code:`info.names`
#
# By this method we are only clearing the AST tree from our class definition
# so it will look like: :code:`class <EnumName>(Enum): pass`
#
# But for example :code:`names` (attributes and methods) for our class will be still defined
# also all inherited things will be still visible.
# This way we can tell mypy that this class has some attributes, it just is not defined in AST
context.cls.defs.body = [
node for node in context.cls.defs.body
if not isinstance(node, nodes.AssignmentStmt)
]
# Create common types handlers
str_type = context.api.named_type('__builtins__.str')
# When working with classes in mypy types
# the only viable option is to use :code:`Instance` even for the class itself
# because even class definition is instance of it's "type".
self_type = types.Instance(info, [])
metaclass_type = info.metaclass_type
# Override __next__ and __getitem__ in the Enum class (for each subclass - each enum)
# to properly say that its return values are children instances.
#
# This is needed as a little hack because we are defining :code:`__iter__` and :code:`__next__`
# in :code:`transform_enum_type` handler using :code:`TypeVar` but
# we are missing something and mypy does not see that when we have a specific implementation
# it should return itself.
# So as a hotfix we redefine these methods on our inherited enum class
# to have more specific typing (which is still not violating original metaclass typing)
_define_method(context, metaclass_type.type,
metaclass_type.type.fullname(), '__next__', [], self_type)
# Example how these all lines below would look like in Python:
#
# ::
# def __getitem__(cls: <self_type>, key: str) -> <self_type>:
# pass
#
#
_define_method(context, metaclass_type.type,
metaclass_type.type.fullname(), '__getitem__', [
nodes.Argument(nodes.Var('cls', self_type), self_type,
None, nodes.ARG_POS),
nodes.Argument(nodes.Var('key', str_type), str_type,
None, nodes.ARG_POS),
], self_type)
# In the end we have to update type of our attributes to return the proper type.
# So we go through all the names defined on the class
# and filter out only type :code:`nodes.Var` which are class attributes.
for name, named_node in info.names.items():
# We want to modify only class an instance level variables
if isinstance(named_node.node, nodes.Var):
node: nodes.Var = info.names[name].node
# We replace original type (which will be int, str, ...)
# with an ``Instance`` of our class itself (not with the base class - Enum).
node.type = types.Instance(info, [])
# We also want to make sure these variables are class-level so you can call
# something like :code:`Color.RED`
node.is_initialized_in_class = True
# In the end assign it back and mark it as generated by plugin
info.names[name] = nodes.SymbolTableNode(nodes.MDEF,
node,
plugin_generated=True)
def from_overloadedfuncdef(
stub: nodes.OverloadedFuncDef) -> "Signature[nodes.Argument]":
"""Returns a Signature from an OverloadedFuncDef.
If life were simple, to verify_overloadedfuncdef, we'd just verify_funcitem for each of its
items. Unfortunately, life isn't simple and overloads are pretty deceitful. So instead, we
try and combine the overload's items into a single signature that is compatible with any
lies it might try to tell.
"""
# For most dunder methods, just assume all args are positional-only
assume_positional_only = is_dunder(stub.name, exclude_special=True)
all_args = {} # type: Dict[str, List[Tuple[nodes.Argument, int]]]
for func in map(_resolve_funcitem_from_decorator, stub.items):
assert func is not None
args = maybe_strip_cls(stub.name, func.arguments)
for index, arg in enumerate(args):
# For positional-only args, we allow overloads to have different names for the same
# argument. To accomplish this, we just make up a fake index-based name.
name = ("__{}".format(index)
if arg.variable.name.startswith("__")
or assume_positional_only else arg.variable.name)
all_args.setdefault(name, []).append((arg, index))
def get_position(arg_name: str) -> int:
# We just need this to return the positional args in the correct order.
return max(index for _, index in all_args[arg_name])
def get_type(arg_name: str) -> mypy.types.ProperType:
with mypy.state.strict_optional_set(True):
all_types = [
arg.variable.type or arg.type_annotation
for arg, _ in all_args[arg_name]
]
return mypy.typeops.make_simplified_union(
[t for t in all_types if t])
def get_kind(arg_name: str) -> int:
kinds = {arg.kind for arg, _ in all_args[arg_name]}
if nodes.ARG_STAR in kinds:
return nodes.ARG_STAR
if nodes.ARG_STAR2 in kinds:
return nodes.ARG_STAR2
# The logic here is based on two tenets:
# 1) If an arg is ever optional (or unspecified), it is optional
# 2) If an arg is ever positional, it is positional
is_opt = (len(all_args[arg_name]) < len(stub.items)
or nodes.ARG_OPT in kinds
or nodes.ARG_NAMED_OPT in kinds)
is_pos = nodes.ARG_OPT in kinds or nodes.ARG_POS in kinds
if is_opt:
return nodes.ARG_OPT if is_pos else nodes.ARG_NAMED_OPT
return nodes.ARG_POS if is_pos else nodes.ARG_NAMED
sig = Signature() # type: Signature[nodes.Argument]
for arg_name in sorted(all_args, key=get_position):
# example_arg_name gives us a real name (in case we had a fake index-based name)
example_arg_name = all_args[arg_name][0][0].variable.name
arg = nodes.Argument(
nodes.Var(example_arg_name, get_type(arg_name)),
type_annotation=None,
initializer=None,
kind=get_kind(arg_name),
)
if arg.kind in (nodes.ARG_POS, nodes.ARG_OPT):
sig.pos.append(arg)
elif arg.kind in (nodes.ARG_NAMED, nodes.ARG_NAMED_OPT):
sig.kwonly[arg.variable.name] = arg
elif arg.kind == nodes.ARG_STAR:
sig.varpos = arg
elif arg.kind == nodes.ARG_STAR2:
sig.varkw = arg
else:
raise AssertionError
return sig
def to_var(self) -> nodes.Var:
return nodes.Var(self.name, self.type)
