Skip to content

Commit f2dd484

Browse files
puddlypuddly
committed
Fix file object lifetimes in command line tools
1 parent 413f98d commit f2dd484

File tree

8 files changed

+132
-36
lines changed

8 files changed

+132
-36
lines changed

‎tests/tools/test_logging.py

Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
1+
import io
12
import logging
23

34
import pytest
45

56
from zigpy_znp.logger import _TRACE
6-
from zigpy_znp.tools.common import setup_parser
7+
from zigpy_znp.tools.common import UnclosableFile, ClosableFileType, setup_parser
78

89

910
@pytest.mark.parametrize(
@@ -27,3 +28,49 @@ def getLevelName(level):
2728
parser.parse_args(["/dev/null"] + ["-v"] * verbosity)
2829

2930
assert logging.getLogger().level == level
31+
32+
33+
def test_command_close_stdout(tmpdir):
34+
parser = setup_parser("Test parser")
35+
parser.add_argument(
36+
"--input",
37+
"-i",
38+
type=ClosableFileType("rb"),
39+
help="Input .bin file",
40+
default="-",
41+
)
42+
43+
parser.add_argument(
44+
"--output",
45+
"-o",
46+
type=ClosableFileType("w"),
47+
help="Output .txt file",
48+
default="-",
49+
)
50+
51+
parser.add_argument(
52+
"--other",
53+
"-t",
54+
type=ClosableFileType("w"),
55+
help="Other .txt file",
56+
required=True,
57+
)
58+
59+
args = parser.parse_args(["/dev/null", "-t", str(tmpdir / "test.txt")])
60+
61+
assert isinstance(args.input, UnclosableFile)
62+
assert isinstance(args.output, UnclosableFile)
63+
assert isinstance(args.other, io.TextIOWrapper)
64+
65+
with args.input as _:
66+
pass
67+
68+
with args.output as _:
69+
pass
70+
71+
with args.other as _:
72+
pass
73+
74+
assert not args.input.closed
75+
assert not args.output.closed
76+
assert args.other.closed

‎zigpy_znp/tools/common.py

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from __future__ import annotations
22

3+
import sys
34
import logging
45
import argparse
56

@@ -142,6 +143,44 @@ def parse_args(self, args: list[str] = None, namespace=None):
142143
return args
143144

144145

146+
class UnclosableFile:
147+
"""
148+
Wraps a file object so that every operation but "close" is proxied.
149+
"""
150+
151+
def __init__(self, f):
152+
self.f = f
153+
154+
def close(self):
155+
return
156+
157+
def __enter__(self):
158+
return
159+
160+
def __exit__(self, exc_type, exc_value, traceback):
161+
return
162+
163+
def __getattr__(self, name):
164+
if name in ("f", "close", "__enter__", "__exit__"):
165+
return super().__getattr__(name)
166+
167+
return getattr(self.f, name)
168+
169+
170+
class ClosableFileType(argparse.FileType):
171+
"""
172+
Allows `FileType` to always be closed properly, even with stdout and stdin.
173+
"""
174+
175+
def __call__(self, string):
176+
f = super().__call__(string)
177+
178+
if f not in (sys.stdin, sys.stdout):
179+
return f
180+
181+
return UnclosableFile(f)
182+
183+
145184
def setup_parser(description: str) -> argparse.ArgumentParser:
146185
"""
147186
Creates an ArgumentParser that sets up a logger with a configurable verbosity

‎zigpy_znp/tools/flash_read.py

Lines changed: 16 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,13 @@
11
import sys
22
import asyncio
33
import logging
4-
import argparse
54

65
import async_timeout
76

87
import zigpy_znp.commands as c
98
from zigpy_znp.api import ZNP
109
from zigpy_znp.config import CONFIG_SCHEMA
11-
from zigpy_znp.tools.common import setup_parser
10+
from zigpy_znp.tools.common import ClosableFileType, setup_parser
1211

1312
LOGGER = logging.getLogger(__name__)
1413

@@ -59,28 +58,32 @@ async def main(argv):
5958
parser.add_argument(
6059
"--output",
6160
"-o",
62-
type=argparse.FileType("wb"),
61+
type=ClosableFileType("wb"),
6362
help="Output .bin file",
6463
required=True,
6564
)
6665

6766
args = parser.parse_args(argv)
6867

69-
znp = ZNP(
70-
CONFIG_SCHEMA(
71-
{"znp_config": {"skip_bootloader": False}, "device": {"path": args.serial}}
68+
with args.output as f:
69+
znp = ZNP(
70+
CONFIG_SCHEMA(
71+
{
72+
"znp_config": {"skip_bootloader": False},
73+
"device": {"path": args.serial},
74+
}
75+
)
7276
)
73-
)
7477

75-
# The bootloader handshake must be the very first command
76-
await znp.connect(test_port=False)
78+
# The bootloader handshake must be the very first command
79+
await znp.connect(test_port=False)
7780

78-
data = await read_firmware(znp)
79-
znp.close()
81+
data = await read_firmware(znp)
82+
znp.close()
8083

81-
args.output.write(data)
84+
f.write(data)
8285

83-
LOGGER.info("Unplug your adapter to leave bootloader mode!")
86+
LOGGER.info("Unplug your adapter to leave bootloader mode!")
8487

8588

8689
if __name__ == "__main__":

‎zigpy_znp/tools/flash_write.py

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,14 @@
33
import sys
44
import asyncio
55
import logging
6-
import argparse
76

87
import async_timeout
98

109
import zigpy_znp.types as t
1110
import zigpy_znp.commands as c
1211
from zigpy_znp.api import ZNP
1312
from zigpy_znp.config import CONFIG_SCHEMA
14-
from zigpy_znp.tools.common import setup_parser
13+
from zigpy_znp.tools.common import ClosableFileType, setup_parser
1514
from zigpy_znp.tools.nvram_reset import nvram_reset
1615

1716
LOGGER = logging.getLogger(__name__)
@@ -143,7 +142,7 @@ async def main(argv):
143142
parser.add_argument(
144143
"--input",
145144
"-i",
146-
type=argparse.FileType("rb"),
145+
type=ClosableFileType("rb"),
147146
help="Input .bin file",
148147
required=True,
149148
)
@@ -157,6 +156,9 @@ async def main(argv):
157156

158157
args = parser.parse_args(argv)
159158

159+
with args.input as f:
160+
firmware = f.read()
161+
160162
znp = ZNP(
161163
CONFIG_SCHEMA(
162164
{"znp_config": {"skip_bootloader": False}, "device": {"path": args.serial}}
@@ -166,7 +168,7 @@ async def main(argv):
166168
# The bootloader handshake must be the very first command
167169
await znp.connect(test_port=False)
168170

169-
await write_firmware(znp=znp, firmware=args.input.read(), reset_nvram=args.reset)
171+
await write_firmware(znp=znp, firmware=firmware, reset_nvram=args.reset)
170172

171173
znp.close()
172174

‎zigpy_znp/tools/network_backup.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -91,13 +91,14 @@ async def main(argv: list[str]) -> None:
9191
)
9292
args = parser.parse_args(argv)
9393

94-
znp = ZNP(ControllerApplication.SCHEMA({"device": {"path": args.serial}}))
95-
await znp.connect()
94+
with args.output as f:
95+
znp = ZNP(ControllerApplication.SCHEMA({"device": {"path": args.serial}}))
96+
await znp.connect()
9697

97-
backup_obj = await backup_network(znp)
98-
znp.close()
98+
backup_obj = await backup_network(znp)
99+
znp.close()
99100

100-
args.output.write(json.dumps(backup_obj, indent=4))
101+
f.write(json.dumps(backup_obj, indent=4))
101102

102103

103104
if __name__ == "__main__":

‎zigpy_znp/tools/network_restore.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,9 @@ async def main(argv: list[str]) -> None:
114114
)
115115
args = parser.parse_args(argv)
116116

117-
backup = json.load(args.input)
117+
with args.input as f:
118+
backup = json.load(f)
119+
118120
validate_backup_json(backup)
119121

120122
await restore_network(

‎zigpy_znp/tools/nvram_read.py

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,13 @@
22
import json
33
import asyncio
44
import logging
5-
import argparse
65

76
import zigpy_znp.types as t
87
from zigpy_znp.api import ZNP
98
from zigpy_znp.config import CONFIG_SCHEMA
109
from zigpy_znp.exceptions import SecurityError, CommandNotRecognized
1110
from zigpy_znp.types.nvids import NWK_NVID_TABLES, ExNvIds, OsalNvIds
12-
from zigpy_znp.tools.common import setup_parser
11+
from zigpy_znp.tools.common import ClosableFileType, setup_parser
1312

1413
LOGGER = logging.getLogger(__name__)
1514

@@ -81,17 +80,19 @@ async def nvram_read(znp: ZNP):
8180
async def main(argv):
8281
parser = setup_parser("Backup a radio's NVRAM")
8382
parser.add_argument(
84-
"--output", "-o", type=argparse.FileType("w"), help="Output file", default="-"
83+
"--output", "-o", type=ClosableFileType("w"), help="Output file", default="-"
8584
)
8685

8786
args = parser.parse_args(argv)
8887

89-
znp = ZNP(CONFIG_SCHEMA({"device": {"path": args.serial}}))
90-
await znp.connect()
88+
with args.output as f:
89+
znp = ZNP(CONFIG_SCHEMA({"device": {"path": args.serial}}))
90+
await znp.connect()
9191

92-
obj = await nvram_read(znp)
93-
args.output.write(json.dumps(obj, indent=4) + "\n")
94-
znp.close()
92+
obj = await nvram_read(znp)
93+
znp.close()
94+
95+
f.write(json.dumps(obj, indent=4) + "\n")
9596

9697

9798
if __name__ == "__main__":

‎zigpy_znp/tools/nvram_write.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,11 @@
22
import json
33
import asyncio
44
import logging
5-
import argparse
65

76
from zigpy_znp.api import ZNP
87
from zigpy_znp.config import CONFIG_SCHEMA
98
from zigpy_znp.types.nvids import ExNvIds, OsalNvIds
10-
from zigpy_znp.tools.common import setup_parser
9+
from zigpy_znp.tools.common import ClosableFileType, setup_parser
1110

1211
LOGGER = logging.getLogger(__name__)
1312

@@ -49,11 +48,13 @@ async def nvram_write(znp: ZNP, backup):
4948
async def main(argv):
5049
parser = setup_parser("Restore a radio's NVRAM from a previous backup")
5150
parser.add_argument(
52-
"--input", "-i", type=argparse.FileType("r"), help="Input file", required=True
51+
"--input", "-i", type=ClosableFileType("r"), help="Input file", required=True
5352
)
5453

5554
args = parser.parse_args(argv)
56-
backup = json.load(args.input)
55+
56+
with args.input as f:
57+
backup = json.load(f)
5758

5859
znp = ZNP(CONFIG_SCHEMA({"device": {"path": args.serial}}))
5960
await znp.connect()

0 commit comments

Comments
 (0)