init

.venv/lib/python3.8/site-packages/mypyc/irbuild/util.py

@@ -0,0 +1,156 @@
+"""Various utilities that don't depend on other modules in mypyc.irbuild."""
+
+from typing import Dict, Any, Union, Optional
+
+from mypy.nodes import (
+    ClassDef, FuncDef, Decorator, OverloadedFuncDef, StrExpr, CallExpr, RefExpr, Expression,
+    IntExpr, FloatExpr, Var, NameExpr, TupleExpr, UnaryExpr, BytesExpr,
+    ArgKind, ARG_NAMED, ARG_NAMED_OPT, ARG_POS, ARG_OPT, GDEF,
+)
+
+
+DATACLASS_DECORATORS = {
+    'dataclasses.dataclass',
+    'attr.s',
+    'attr.attrs',
+}
+
+
+def is_trait_decorator(d: Expression) -> bool:
+    return isinstance(d, RefExpr) and d.fullname == 'mypy_extensions.trait'
+
+
+def is_trait(cdef: ClassDef) -> bool:
+    return any(is_trait_decorator(d) for d in cdef.decorators) or cdef.info.is_protocol
+
+
+def dataclass_decorator_type(d: Expression) -> Optional[str]:
+    if isinstance(d, RefExpr) and d.fullname in DATACLASS_DECORATORS:
+        return d.fullname.split('.')[0]
+    elif (isinstance(d, CallExpr)
+            and isinstance(d.callee, RefExpr)
+            and d.callee.fullname in DATACLASS_DECORATORS):
+        name = d.callee.fullname.split('.')[0]
+        if name == 'attr' and 'auto_attribs' in d.arg_names:
+            # Note: the mypy attrs plugin checks that the value of auto_attribs is
+            # not computed at runtime, so we don't need to perform that check here
+            auto = d.args[d.arg_names.index('auto_attribs')]
+            if isinstance(auto, NameExpr) and auto.name == 'True':
+                return 'attr-auto'
+        return name
+    else:
+        return None
+
+
+def is_dataclass_decorator(d: Expression) -> bool:
+    return dataclass_decorator_type(d) is not None
+
+
+def is_dataclass(cdef: ClassDef) -> bool:
+    return any(is_dataclass_decorator(d) for d in cdef.decorators)
+
+
+def dataclass_type(cdef: ClassDef) -> Optional[str]:
+    for d in cdef.decorators:
+        typ = dataclass_decorator_type(d)
+        if typ is not None:
+            return typ
+    return None
+
+
+def get_mypyc_attr_literal(e: Expression) -> Any:
+    """Convert an expression from a mypyc_attr decorator to a value.
+
+    Supports a pretty limited range."""
+    if isinstance(e, (StrExpr, IntExpr, FloatExpr)):
+        return e.value
+    elif isinstance(e, RefExpr) and e.fullname == 'builtins.True':
+        return True
+    elif isinstance(e, RefExpr) and e.fullname == 'builtins.False':
+        return False
+    elif isinstance(e, RefExpr) and e.fullname == 'builtins.None':
+        return None
+    return NotImplemented
+
+
+def get_mypyc_attr_call(d: Expression) -> Optional[CallExpr]:
+    """Check if an expression is a call to mypyc_attr and return it if so."""
+    if (
+        isinstance(d, CallExpr)
+        and isinstance(d.callee, RefExpr)
+        and d.callee.fullname == 'mypy_extensions.mypyc_attr'
+    ):
+        return d
+    return None
+
+
+def get_mypyc_attrs(stmt: Union[ClassDef, Decorator]) -> Dict[str, Any]:
+    """Collect all the mypyc_attr attributes on a class definition or a function."""
+    attrs: Dict[str, Any] = {}
+    for dec in stmt.decorators:
+        d = get_mypyc_attr_call(dec)
+        if d:
+            for name, arg in zip(d.arg_names, d.args):
+                if name is None:
+                    if isinstance(arg, StrExpr):
+                        attrs[arg.value] = True
+                else:
+                    attrs[name] = get_mypyc_attr_literal(arg)
+
+    return attrs
+
+
+def is_extension_class(cdef: ClassDef) -> bool:
+    if any(
+        not is_trait_decorator(d)
+        and not is_dataclass_decorator(d)
+        and not get_mypyc_attr_call(d)
+        for d in cdef.decorators
+    ):
+        return False
+    if cdef.info.typeddict_type:
+        return False
+    if cdef.info.is_named_tuple:
+        return False
+    if (cdef.info.metaclass_type and cdef.info.metaclass_type.type.fullname not in (
+            'abc.ABCMeta', 'typing.TypingMeta', 'typing.GenericMeta')):
+        return False
+    return True
+
+
+def get_func_def(op: Union[FuncDef, Decorator, OverloadedFuncDef]) -> FuncDef:
+    if isinstance(op, OverloadedFuncDef):
+        assert op.impl
+        op = op.impl
+    if isinstance(op, Decorator):
+        op = op.func
+    return op
+
+
+def concrete_arg_kind(kind: ArgKind) -> ArgKind:
+    """Find the concrete version of an arg kind that is being passed."""
+    if kind == ARG_OPT:
+        return ARG_POS
+    elif kind == ARG_NAMED_OPT:
+        return ARG_NAMED
+    else:
+        return kind
+
+
+def is_constant(e: Expression) -> bool:
+    """Check whether we allow an expression to appear as a default value.
+
+    We don't currently properly support storing the evaluated
+    values for default arguments and default attribute values, so
+    we restrict what expressions we allow.  We allow literals of
+    primitives types, None, and references to Final global
+    variables.
+    """
+    return (isinstance(e, (StrExpr, BytesExpr, IntExpr, FloatExpr))
+            or (isinstance(e, UnaryExpr) and e.op == '-'
+                and isinstance(e.expr, (IntExpr, FloatExpr)))
+            or (isinstance(e, TupleExpr)
+                and all(is_constant(e) for e in e.items))
+            or (isinstance(e, RefExpr) and e.kind == GDEF
+                and (e.fullname in ('builtins.True', 'builtins.False', 'builtins.None')
+                     or (isinstance(e.node, Var) and e.node.is_final))))