0

I have a custom test runner that uses inspect to find all the functions starting with with test_. At one point I add an attribute to the function object. This is done so the test runner has access to the docstring information later.

I'd really like to create a type hint for the list of functions so that my IDE (pycharm) knows that tdi is indeed an attribute of the function object.

How do you build such a type hint?

for m, fn in module.__dict__.items():
        if m.startswith(prefix):
            if callable(fn):
                tdi = getTestDocInfo(fn) # reads the docstring for tags and puts the info in a dataclass
                fn.tdi = tdi
                tests.append(fn)

edit: In this case, the type hint should convey the object is a function AND that it has the attribute tdi.

As @chepner pointed out the Protocol building block for a type hint is what was needed here.

This is what worked.

@runtime_checkable
class TestFunc(Protocol):
    tdi: TestDocInfo
    def __call__(self, testresult: TestResult, shareditems: SharedItems): pass

tests: List[TestFunc]. #type hint I am looking for.
Marcel Wilson
  • 3,842
  • 1
  • 26
  • 55
  • Looks like a duplicate of [Python type hint for classes that support __getitem__](https://stackoverflow.com/q/55127855/7851470) (the answer there should be updated, though). See also [How can I use static checking to ensure an object has a certain method/attribute?](https://stackoverflow.com/questions/64185810/how-can-i-use-static-checking-to-ensure-an-object-has-a-certain-method-attribute). – Georgy Nov 11 '20 at 22:47
  • Does this answer your question? [Python type hint for classes that support \_\_getitem\_\_](https://stackoverflow.com/questions/55127855/python-type-hint-for-classes-that-support-getitem) – Georgy Nov 12 '20 at 22:55
  • 1
    Initially I didn't think those were related due to the nature of adding custom attributes to a `function` object. But I think you're right, these are essentially the same issue. The only real difference is I still needed the type hint to convey the fact it is also a `function`. Which required me to include a definition of `__call__` as well. – Marcel Wilson Nov 13 '20 at 18:05

1 Answers1

1

This seems like a job for typing.Protocol and typing.runtime_checkable:

from typing import Protocol, runtime_checkable

@runtime_checkable
class HasTDI(Protocol):
    tdi: int  # Or whatever type is appropriate.

Then

isinstance(fn, HasTDI)

should be true if fn.tdi exists and is an int.

I'm afraid I don't know if PyCharm will make use of this information, though.

chepner
  • 497,756
  • 71
  • 530
  • 681
  • This is really useful. Though I may have to revise my question. I'm curious if I can still retain the the type hint of 'FunctionType' or 'Callable'. – Marcel Wilson Nov 11 '20 at 23:04
  • Oh derp, `List[Union[types.FunctionType, HasTDI]]` works like a charm! – Marcel Wilson Nov 11 '20 at 23:06
  • `Union` implies that an item could be a function without a `tdii` attribute, or something other than a function (or callable in general) that does. Maybe you want `HasTDI` to subclass `FunctionType` as well instead? E.g., `class HasTDI(Protocol, FunctionType)` or something similar, then `List[HasTDI]`. – chepner Nov 12 '20 at 01:08
  • You're right. I played with it afterward and found the union isn't quite right. I ended up just adding `def __call__(self, testresult: TestResult, shareditems: SharedItems): pass` to the `HasTDI` then just using `List[HasTDI]`. Thanks again for the heads up on `Protocol` – Marcel Wilson Nov 12 '20 at 14:10