[Python] Incompatible type hints for AST node subclasses

I am writing a script to perform obfuscation using the ast module. A decorator is used to define obfuscation functions for each AST node type. The decorator adds the functions to a dictionary, Obfuscator.obfuscators, which maps type(node) to a function that takes the node and context and returns a node with the same parent (ast.Name, for example, is a subclass of ast.expr and its obfuscator function can return any subclass of ast.expr from context).

import ast
from typing import Callable, TypeVar, cast, overload

Nullable = ast.stmt
NonNullable = ast.expr | ast.comprehension | ast.keyword
NullableT = TypeVar("NullableT", bound=Nullable)
NonNullableT = TypeVar("NonNullableT", bound=NonNullable)

T = TypeVar("T", bound=Nullable | NonNullable)
A = TypeVar("A", bound=ast.AST)

Context = dict[str, ast.expr]


class Obfuscator:
obfuscators: dict[type[ast.AST], Callable[[ast.AST, Context], ast.AST | None]] = {}

@classmethod
@overload
def obfuscate(cls, *nodes: NonNullableT, context: Context = {}) -> NonNullableT: ...

@classmethod
@overload
def obfuscate(
cls, *nodes: NullableT, context: Context = {}
) -> NullableT | None: ...

@classmethod
def obfuscate(cls, *nodes: T, context: Context = {}) -> T | None:
context = dict(context)
for node in nodes:
obfuscator = cls.obfuscators.get(type(node))
if obfuscator:
result = obfuscator(node, context)
if result:
return cast(T, result)
else:
return node

@classmethod
def obfuscate_all(cls, nodes: list[T], context: Context) -> list[T]:
return [
x
for x in [Obfuscator.obfuscate(node, context=context) for node in nodes]
if x is not None
]

@classmethod
def define(
cls,
ast_type: type[A],
):
def decorator(func: Callable[[A, Context], T | None]):
# this is the line that causes the pyright error
cls.obfuscators[ast_type] = func
return func

return decorator


Some example definitions:

@Obfuscator.define(ast.Import)
def obf_import(stmt: ast.Import, context: Context):
# replace `import x` with `__import__(x)`
for name in stmt.names:
context[name.asname or name.name] = ast.Call(
ast.Name(id="__import__"), args=[ast.Constant(name.name)], keywords=[]
)


@Obfuscator.define(ast.With)
def obf_with(stmt: ast.With, context: Context):
for with_ in stmt.items:
as_var: ast.expr | None = with_.optional_vars
if isinstance(as_var, ast.Name) and as_var.id:
context[as_var.id] = Obfuscator.obfuscate(
with_.context_expr, context=context
)
return Obfuscator.obfuscate(*stmt.body, context=context)


@Obfuscator.define(ast.GeneratorExp)
def obf_generatorexp(expr: ast.GeneratorExp, context: Context):
expr.elt = Obfuscator.obfuscate(expr.elt, context=context)
expr.generators = Obfuscator.obfuscate_all(expr.generators, context=context)
return expr


@Obfuscator.define(ast.Name)
def obf_name(expr: ast.Name, context: Context):
return context.get(expr.id, expr)


This code works fine, but I cannot work out how to type hint define. Pyright currently gives one error:

Argument of type "(A@define, Context) -> (T@decorator | None)" cannot be assigned to parameter "value" of type "(AST, Context) -> (AST | None)" in function "__setitem__"
Type "(A@define, Context) -> (T@decorator | None)" is not assignable to type "(AST, Context) -> (AST | None)"
Parameter 1: type "AST" is incompatible with type "A@define"
Type "AST" is not assignable to type "A@define"


I think this is because the type hint for obfuscators means the functions are expected to work for any AST node, but they only work for one.

What is the best way to type hint define and obfuscators?

Continue reading...