Extending NewType to support "spreading"

[ddworken]

Currently NewType doesn't have any way of "spreading" whereby UserID(1) + UserID(2) would be a UserID (instead it is an int). I'm working on implementing a form of taint analysis using mypy for compile time detection of tainted data hitting sensitive sinks and adding this would be unbelievably helpful for my project and I imagine many others. Additionally, this was actually one of the suggested uses of NewType in #1284 but as far as I can tell it isn't possible given the current implementation of NewType.

The ultimate goal would be to make it so that something like this would pass type checking:

from typing import NewType, TypeVar, Callable

Type = NewType('Type', str, spread=True)

def src() -> Type:
	return Type("Type!")

def dst(data: Type) -> None:
	print(data)

dst(src())  # Currently passes
dst(src()+"OtherData")  # Currently fails type checking
dst("OtherData"+src())  # Currently fails type checking

What do people think of this as an extension to NewType?

[ddfisher]

I'm not convinced this is generally useful enough to be worth adding to mypy. Instead, you could implement something like this yourself:

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    class TaintedString:
        def __init__(self, x: str) -> None: ...
        def __add__(self, x: str) -> TaintedString: ...
        def __radd__(self, x: str) -> TaintedString: ...
else:
    TaintedString = lambda x: x


def src() -> TaintedString:
	return TaintedString("Type!")

def dst(data: TaintedString) -> None:
	print(data)

dst(src())  # Currently passes
dst(src()+"OtherData")  # Passes
dst("OtherData"+src())  # Passes

One caveat is that TaintedString will need to re-specify all the string methods you want to use with appropriate type signatures, but in the worst case (if you want all of them) you can copy them from typeshed with minor modifications.

[ddworken]

Thanks for the feedback!