Proper typing for sys.stdout and files

Viewed 687

I'm trying to create a method that can handle a string representing a file name, a Path object or an already open output stream (sys.stdout, open('...', 'w')). And I'm trying to properly set the types for mypy to check them.

So far I got:

import io
from pathlib import Path
from typing import Union, TextIO, Text


def generate(output: Union[Text, Path, TextIO]) -> None:
    if isinstance(output, io.IOBase):
        output.write("data")
    else:
        if isinstance(output, Text):
            output = Path(output)
        with output.open("w") as output_file:
            output_file.write("data")

But mypy keeps complaining with

Item "TextIO" of "Union[Path, TextIO]" has no attribute "open"

AFAIK, TextIO is the proper type to use for a text file, but the isinstance check cannot be done against this type. The code structure ensures that at the point of error we cannot have a TestIO object, as it had been processed in the previous branch.

How should I tag all the types here?

1 Answers

This code could be rewritten as follows to isolate open in the Path branch, because Mypy understands isinstance checks, and call write in one place:

import io
from pathlib import Path
from typing import Union, TextIO, Text
from contextlib import ExitStack


def generate(output: Union[Text, Path, TextIO]) -> None:
    with ExitStack() as stack:
        if isinstance(output, Text):
            output = Path(output)
        if isinstance(output, Path):
            output = stack.enter_context(output.open("w"))
        output.write("data")
Related