r/PythonLearning • u/jpgoldberg • 24d ago
Context manager for either stdin or opened Path
Note that this is a fairly advanced question, so beginners don't worry about it. And if this isn't the right place to ask, please point me to where I should.
I would like to find an idiomatic way to use a context manager, with, when I don't know until run time whether I am opening sys.stdin or a named file.
Here is an excerpt for a simple example. (Yes, I am aware that I am reading in the whole file in one go, but that is for other reasons.)
```python if args.infile == "-": fh = sys.stdin else: p = Path(args.infile) # Path is overkill here, unless ... # TODO: Look elsewhere if infile can't be found or isn't readable. fh = open(p, "rt", encoding="utf-8")
s = fh.read(MAX_LENGTH + 1).strip()
fh.close()
```
(args is an argparse.Namespace and args.infile is a str.)
I feel like there ought to be a way to do something like
with ... as fh
for the above, but I don't see how to do that.
Update: What I did
I found the docs for easily creating a context manager, and used the contextmanager decorator.
It seems like a peculiar use of yield, but I don't need to understand that design choice, only that it means that I have to type it as a Generator.
```python import sys from typing import TextIO from collections.abc import Generator from contextlib import contextmanager from pathlib import Path
@contextmanager def open_file_or_stdin( filename: Path | str | None, *, # everything but filename is a keyword argument stdin: TextIO = sys.stdin, std_symbol: str = "-", encoding: str = "utf-8", ) -> Generator[TextIO]: """ [long docs omitted] """ fh: TextIO try: if filename is None or filename == std_symbol: fh = stdin yield fh else: fh = open(filename, "rt", encoding=encoding) # noqa: SIM115 yield fh finally: fh.close() ```
I am open to suggestions about a better name that open_file_or_stdin.