Skip to content

Commit e8a4d80

Browse files
1e0ng1e0ng
committed
Add a new download() method for PybricksHub.
1 parent 399ba43 commit e8a4d80

File tree

3 files changed

+77
-90
lines changed

3 files changed

+77
-90
lines changed

‎pybricksdev/cli/__init__.py‎

Lines changed: 1 addition & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -325,21 +325,7 @@ def is_pybricks_usb(dev):
325325
await hub.connect()
326326
try:
327327
with _get_script_path(args.file) as script_path:
328-
# For PybricksHub, we use download_user_program to just upload without running
329-
if args.conntype in ["ble", "usb"]:
330-
from pybricksdev.compile import compile_multi_file
331-
332-
# Compile the script to mpy format
333-
mpy = await compile_multi_file(
334-
script_path, hub._mpy_abi_version or 6
335-
)
336-
# Upload without running
337-
await hub.download_user_program(mpy)
338-
# For EV3Connection, we just use download
339-
elif args.conntype == "ssh":
340-
await hub.download(script_path)
341-
else:
342-
raise RuntimeError("Unsupported hub type for download command")
328+
await hub.download(script_path)
343329
finally:
344330
await hub.disconnect()
345331

‎pybricksdev/connections/pybricks.py‎

Lines changed: 47 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -466,6 +466,50 @@ async def stop_user_program(self) -> None:
466466
response=True,
467467
)
468468

469+
async def download(self, script_path: str) -> None:
470+
"""
471+
Downloads a script to the hub without running it.
472+
473+
This method handles both compilation and downloading of the script.
474+
For Pybricks hubs, it compiles the script to MPY format and downloads it
475+
using the Pybricks protocol.
476+
477+
Args:
478+
script_path: Path to the Python script to download.
479+
480+
Raises:
481+
RuntimeError: If the hub is not connected or if the hub type is not supported.
482+
ValueError: If the compiled program is too large to fit on the hub.
483+
"""
484+
if self.connection_state_observable.value != ConnectionState.CONNECTED:
485+
raise RuntimeError("not connected")
486+
487+
# since Pybricks profile v1.2.0, the hub will tell us which file format(s) it supports
488+
if not (
489+
self._capability_flags
490+
& (
491+
HubCapabilityFlag.USER_PROG_MULTI_FILE_MPY6
492+
| HubCapabilityFlag.USER_PROG_MULTI_FILE_MPY6_1_NATIVE
493+
)
494+
):
495+
raise RuntimeError(
496+
"Hub is not compatible with any of the supported file formats"
497+
)
498+
499+
# no support for native modules unless one of the flags below is set
500+
abi = 6
501+
502+
if (
503+
self._capability_flags
504+
& HubCapabilityFlag.USER_PROG_MULTI_FILE_MPY6_1_NATIVE
505+
):
506+
abi = (6, 1)
507+
508+
# Compile the script to mpy format
509+
mpy = await compile_multi_file(script_path, abi)
510+
# Download without running
511+
await self.download_user_program(mpy)
512+
469513
async def run(
470514
self,
471515
py_path: Optional[str] = None,
@@ -506,31 +550,11 @@ async def run(
506550
await self._legacy_run(py_path, wait)
507551
return
508552

509-
# since Pybricks profile v1.2.0, the hub will tell us which file format(s) it supports
510-
if not (
511-
self._capability_flags
512-
& (
513-
HubCapabilityFlag.USER_PROG_MULTI_FILE_MPY6
514-
| HubCapabilityFlag.USER_PROG_MULTI_FILE_MPY6_1_NATIVE
515-
)
516-
):
517-
raise RuntimeError(
518-
"Hub is not compatible with any of the supported file formats"
519-
)
520-
521-
# no support for native modules unless one of the flags below is set
522-
abi = 6
523-
524-
if (
525-
self._capability_flags
526-
& HubCapabilityFlag.USER_PROG_MULTI_FILE_MPY6_1_NATIVE
527-
):
528-
abi = (6, 1)
529-
553+
# Download the program if a path is provided
530554
if py_path is not None:
531-
mpy = await compile_multi_file(py_path, abi)
532-
await self.download_user_program(mpy)
555+
await self.download(py_path)
533556

557+
# Start the program
534558
await self.start_user_program()
535559

536560
if wait:

‎tests/test_cli.py‎

Lines changed: 29 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ async def test_download_ble(self):
6464
# Create a mock hub
6565
mock_hub = AsyncMock()
6666
mock_hub._mpy_abi_version = 6
67-
mock_hub.download_user_program = AsyncMock()
67+
mock_hub.download = AsyncMock()
6868

6969
# Create a temporary file
7070
with tempfile.NamedTemporaryFile(suffix=".py", mode="w+", delete=False) as temp:
@@ -79,27 +79,20 @@ async def test_download_ble(self):
7979
name="MyHub",
8080
)
8181

82-
# Mock the compilation
83-
mock_mpy = MagicMock()
84-
mock_mpy.__bytes__ = lambda self: b"compiled code"
85-
8682
# Mock the hub creation
8783
with patch(
8884
"pybricksdev.connections.pybricks.PybricksHubBLE", return_value=mock_hub
8985
) as mock_hub_class:
9086
with patch("pybricksdev.ble.find_device", return_value="mock_device"):
91-
with patch(
92-
"pybricksdev.compile.compile_multi_file", return_value=mock_mpy
93-
):
94-
# Run the command
95-
download = Download()
96-
await download.run(args)
87+
# Run the command
88+
download = Download()
89+
await download.run(args)
9790

98-
# Verify the hub was created and used correctly
99-
mock_hub_class.assert_called_once_with("mock_device")
100-
mock_hub.connect.assert_called_once()
101-
mock_hub.download_user_program.assert_called_once()
102-
mock_hub.disconnect.assert_called_once()
91+
# Verify the hub was created and used correctly
92+
mock_hub_class.assert_called_once_with("mock_device")
93+
mock_hub.connect.assert_called_once()
94+
mock_hub.download.assert_called_once()
95+
mock_hub.disconnect.assert_called_once()
10396
finally:
10497
os.unlink(temp_path)
10598

@@ -109,7 +102,7 @@ async def test_download_usb(self):
109102
# Create a mock hub
110103
mock_hub = AsyncMock()
111104
mock_hub._mpy_abi_version = 6
112-
mock_hub.download_user_program = AsyncMock()
105+
mock_hub.download = AsyncMock()
113106

114107
# Create a temporary file
115108
with tempfile.NamedTemporaryFile(suffix=".py", mode="w+", delete=False) as temp:
@@ -124,27 +117,20 @@ async def test_download_usb(self):
124117
name=None,
125118
)
126119

127-
# Mock the compilation
128-
mock_mpy = MagicMock()
129-
mock_mpy.__bytes__ = lambda self: b"compiled code"
130-
131120
# Mock the hub creation
132121
with patch(
133122
"pybricksdev.connections.pybricks.PybricksHubUSB", return_value=mock_hub
134123
) as mock_hub_class:
135124
with patch("usb.core.find", return_value="mock_device"):
136-
with patch(
137-
"pybricksdev.compile.compile_multi_file", return_value=mock_mpy
138-
):
139-
# Run the command
140-
download = Download()
141-
await download.run(args)
125+
# Run the command
126+
download = Download()
127+
await download.run(args)
142128

143-
# Verify the hub was created and used correctly
144-
mock_hub_class.assert_called_once_with("mock_device")
145-
mock_hub.connect.assert_called_once()
146-
mock_hub.download_user_program.assert_called_once()
147-
mock_hub.disconnect.assert_called_once()
129+
# Verify the hub was created and used correctly
130+
mock_hub_class.assert_called_once_with("mock_device")
131+
mock_hub.connect.assert_called_once()
132+
mock_hub.download.assert_called_once()
133+
mock_hub.disconnect.assert_called_once()
148134
finally:
149135
os.unlink(temp_path)
150136

@@ -214,7 +200,7 @@ async def test_download_stdin(self):
214200
# Create a mock hub
215201
mock_hub = AsyncMock()
216202
mock_hub._mpy_abi_version = 6
217-
mock_hub.download_user_program = AsyncMock()
203+
mock_hub.download = AsyncMock()
218204

219205
# Create a mock stdin
220206
mock_stdin = io.StringIO("print('test')")
@@ -228,31 +214,22 @@ async def test_download_stdin(self):
228214
name="MyHub",
229215
)
230216

231-
# Mock the compilation
232-
mock_mpy = MagicMock()
233-
mock_mpy.__bytes__ = lambda self: b"compiled code"
234-
235217
# Mock the hub creation and file handling
236218
with patch(
237219
"pybricksdev.connections.pybricks.PybricksHubBLE", return_value=mock_hub
238220
) as mock_hub_class:
239221
with patch("pybricksdev.ble.find_device", return_value="mock_device"):
240-
with patch(
241-
"pybricksdev.compile.compile_multi_file", return_value=mock_mpy
242-
):
243-
with patch("tempfile.NamedTemporaryFile") as mock_temp:
244-
mock_temp.return_value.__enter__.return_value.name = (
245-
"/tmp/test.py"
246-
)
247-
# Run the command
248-
download = Download()
249-
await download.run(args)
222+
with patch("tempfile.NamedTemporaryFile") as mock_temp:
223+
mock_temp.return_value.__enter__.return_value.name = "/tmp/test.py"
224+
# Run the command
225+
download = Download()
226+
await download.run(args)
250227

251-
# Verify the hub was created and used correctly
252-
mock_hub_class.assert_called_once_with("mock_device")
253-
mock_hub.connect.assert_called_once()
254-
mock_hub.download_user_program.assert_called_once()
255-
mock_hub.disconnect.assert_called_once()
228+
# Verify the hub was created and used correctly
229+
mock_hub_class.assert_called_once_with("mock_device")
230+
mock_hub.connect.assert_called_once()
231+
mock_hub.download.assert_called_once()
232+
mock_hub.disconnect.assert_called_once()
256233

257234
@pytest.mark.asyncio
258235
async def test_download_connection_error(self):

0 commit comments

Comments
 (0)