Type Tailoring without Macros

The dream would be to build an ideally flexible system that can typecheck any dynamically-typed program. But we are not there yet - and mostly because of the non-disjoint nature of data types in these languages. Dynamic languages allow types to overlap and inherit from each other (for example, in Python a value can be both a string and an iterable), making it nearly impossible for type checkers to reason about all possible type relationships¹. Until we build these infallible type checkers, maybe we can explore developer-assisted approaches to nudge our existing systems.

One such approach is type tailoring². Type tailoring is a way to extend traditional type systems with domain-specific types. While intended for languages with macro systems and runtime checks, the same goal could be realized with extensible type checkers with a plugin interface such as Mypy³ that allow developers to define constructs for domain-specific (or third-party) types or overlay the type system with domain-specific typing rules.

This piece highlights an interesting use case of Mypy plugins: Lazy evaluation inference, and discusses an idealistic design for enforcing custom type-checking behaviour (by preprocessing Type Hints).

Lazy evaluation inference

The Lazy evaluation pattern is commonplace within the python community for delaying the instantiation of objects until they are needed. This approach optimizes memory by creating objects only at the point of first access, rather than at module import time. But because type information about this lazy object is not available until runtime, type checkers are unable to determine types during static analysis. Idioms like these present challenges for gradual type checkers like Mypy, often forcing developers to fall back to the Any type, which weakens soundness guarantees.

Consider this code sample of lazy object instantiation:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

from typing import Any, TypeVar, Generic, Callable, Self

empty = object()
Obj = TypeVar('Obj')

class Lazy(Generic[Obj]):
	def __init__(self: Self, f: Callable[[], Obj]) -> None:
		self._f = f
		self._inst = empty
	
	def __getattr__(self: Self, attr: str) -> Obj:
	if self._inst == empty:
		self._inst = self._f()
		
	return getattr(self._inst, attr)

class NewObj:
	def __init__(self: Self) -> None:
		print("finally created")
	
	def foo(self: Self) -> None:
		print("lookey here. foo() called")
	
def make_obj() -> NewObj:
	return NewObj()

lazy_inst = Lazy(make_obj)

lazy_inst.foo()

Mypy fails to recognize that foo() is a method of our lazily instantiated object. By running reveal_type(lazy_inst.foo()), we see Mypy resorts to Any instead of capturing None as the return type effectively losing all type safety for operations on our lazy object.

1
2
3
4
5

reveal_type(lazy_inst) 
# note: Revealed type is "Lazy[NewObj]"

reveal_type(lazy_inst.foo()) 
# note: Revealed type is "Any"

This means we lose precision about our lazy object and further interactions with the object just propagate the Any type all up our type analysis.

This is where Mypy plugins come in. By implementing attribute hooks, we can teach Mypy how to correctly type lazy objects. These hooks intercept attribute lookups during static analysis, allowing us to provide Mypy with the actual types that will be available at runtime. For our lazy object, we’ll create a plugin that tells Mypy that any attribute access should resolve to the types of the underlying NewObj instance.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

from functools import partial
from mypy.errorcodes import ATTR_DEFINED
from mypy.plugin import Plugin, AttributeContext
from mypy.subtypes import find_member
from typing import Optional, Callable, Any
from mypy.types import Type, AnyType, Instance

def lazy_plugin_attribute_hook(ctx: AttributeContext, *, attr: str) -> Type:
	if not isinstance(ctx.default_attr_type, AnyType):
		return ctx.default_attr_type

	assert len(ctx.type.args) == 1, ctx.type
	gen_type = ctx.type.args[0]
	mem = find_member(attr, gen_type, gen_type)

	if mem:
		return mem
	else:
		ctx.api.fail(f"unknown attr accessed: {attr}", ctx.context,)
	return ctx.default_attr_type

class LazyPlugin(Plugin):
	def get_attribute_hook(self, fullname: str) -> Optional[Callable[[AttributeContext], Type]]:

	if fullname.startswith('projs.test_mypy.lazy.Lazy'):
		_, _, attrs = fullname.rpartition('.')
		return partial(lazy_plugin_attribute_hook, attr=attrs)

	return None

def plugin(version: str) -> type[LazyPlugin]:
	return LazyPlugin

When we call reveal_type again, Mypy correctly identifies lazy_inst.foo() as returning None, giving us proper type safety for our lazy instantiation pattern.

1
2
3
4
5
6
7
8

reveal_type(lazy_inst) 
# note: Revealed type is "Lazy[NewObj]"

reveal_type(lazy_inst.foo()) 
# note: Revealed type is "None"

reveal_type(lazy_inst.not_exists()) 
# note:  error: unknown attr accessed: not_exists [misc]

Preprocessing Type Hints

I explored extending this use of Mypy plugins to attempt preprocess type hints dynamically. For instance, consider a poorly typed function returning Any instead of a more specific type like str.

1
2
3
4
5
6
7

# poorly typed function
def bad(x: int) -> Any:
	return str(x)

# ideal (internal) modification 
def bad(x: int) -> str:
	return str(x)

Could a plugin replace Any with the correct type and enforce type-checking on the updated hints? While I successfully replaced Any with concrete types⁴, I couldn’t control when plugins were called during Mypy’s analysis, making it impossible to enforce typechecking on the updated annotations. Hopefully, future Mypy updates will address these limitations.