Skip to content

Commit 7a1ca00

Browse files
shaggysashaggysa
committed
cli/__init__.py: move stay_connected_menu to separate function
1 parent aedee30 commit 7a1ca00

File tree

1 file changed

+154
-120
lines changed

1 file changed

+154
-120
lines changed

‎pybricksdev/cli/__init__.py‎

Lines changed: 154 additions & 120 deletions
Original file line numberDiff line numberDiff line change
@@ -183,142 +183,174 @@ def add_parser(self, subparsers: argparse._SubParsersAction):
183183
default=False,
184184
)
185185

186-
async def run(self, args: argparse.Namespace):
187-
188-
async def stay_connected_menu(hub: PybricksHub):
186+
async def stay_connected_menu(self, hub: PybricksHub, args: argparse.Namespace):
189187

190-
class ResponseOptions(IntEnum):
191-
RECOMPILE_RUN = 0
192-
RECOMPILE_DOWNLOAD = 1
193-
RUN_STORED = 2
194-
CHANGE_TARGET_FILE = 3
195-
EXIT = 4
188+
if args.conntype == "ble":
189+
from pybricksdev.ble import find_device as find_ble
190+
from pybricksdev.connections.pybricks import PybricksHubBLE
191+
else:
192+
from usb.core import find as find_usb
196193

197-
async def reconnect_hub():
198-
if not await questionary.confirm(
199-
"\nThe hub has been disconnected. Would you like to re-connect?"
200-
).ask_async():
201-
exit()
194+
from pybricksdev.connections.pybricks import PybricksHubUSB
195+
from pybricksdev.usb import (
196+
EV3_USB_PID,
197+
LEGO_USB_VID,
198+
MINDSTORMS_INVENTOR_USB_PID,
199+
NXT_USB_PID,
200+
SPIKE_ESSENTIAL_USB_PID,
201+
SPIKE_PRIME_USB_PID,
202+
)
202203

203-
if args.conntype == "ble":
204-
print(
205-
f"Searching for {args.name or 'any hub with Pybricks service'}..."
204+
def is_pybricks_usb(dev):
205+
return (
206+
(dev.idVendor == LEGO_USB_VID)
207+
and (
208+
dev.idProduct
209+
in [
210+
NXT_USB_PID,
211+
EV3_USB_PID,
212+
SPIKE_PRIME_USB_PID,
213+
SPIKE_ESSENTIAL_USB_PID,
214+
MINDSTORMS_INVENTOR_USB_PID,
215+
]
206216
)
207-
device_or_address = await find_ble(args.name)
208-
hub = PybricksHubBLE(device_or_address)
209-
elif args.conntype == "usb":
210-
device_or_address = find_usb(custom_match=is_pybricks_usb)
211-
hub = PybricksHubUSB(device_or_address)
212-
213-
await hub.connect()
214-
# re-enable echoing of the hub's stdout
215-
hub._enable_line_handler = True
216-
hub.print_output = True
217-
return hub
218-
219-
response_options = [
220-
"Recompile and Run",
221-
"Recompile and Download",
222-
"Run Stored Program",
223-
"Change Target File",
224-
"Exit",
225-
]
226-
# the entry that is selected by default when the menu opens
227-
# this is overridden after the user picks an option
228-
# so that the default option is always the one that was last chosen
229-
default_response_option = (
230-
ResponseOptions.RECOMPILE_RUN
231-
if args.start
232-
else ResponseOptions.RECOMPILE_DOWNLOAD
233-
)
217+
and dev.product.endswith("Pybricks")
218+
)
234219

235-
while True:
236-
try:
237-
if args.file is sys.stdin:
238-
await hub.race_disconnect(
239-
hub.race_power_button_press(
240-
questionary.press_any_key_to_continue(
241-
"The hub will stay connected and echo its output to the terminal. Press any key to exit."
242-
).ask_async()
243-
)
244-
)
245-
return
246-
response = await hub.race_disconnect(
220+
class ResponseOptions(IntEnum):
221+
RECOMPILE_RUN = 0
222+
RECOMPILE_DOWNLOAD = 1
223+
RUN_STORED = 2
224+
CHANGE_TARGET_FILE = 3
225+
EXIT = 4
226+
227+
async def reconnect_hub():
228+
if not await questionary.confirm(
229+
"\nThe hub has been disconnected. Would you like to re-connect?"
230+
).ask_async():
231+
exit()
232+
233+
if args.conntype == "ble":
234+
print(
235+
f"Searching for {args.name or 'any hub with Pybricks service'}..."
236+
)
237+
device_or_address = await find_ble(args.name)
238+
hub = PybricksHubBLE(device_or_address)
239+
elif args.conntype == "usb":
240+
device_or_address = find_usb(custom_match=is_pybricks_usb)
241+
hub = PybricksHubUSB(device_or_address)
242+
243+
await hub.connect()
244+
# re-enable echoing of the hub's stdout
245+
hub._enable_line_handler = True
246+
hub.print_output = True
247+
return hub
248+
249+
response_options = [
250+
"Recompile and Run",
251+
"Recompile and Download",
252+
"Run Stored Program",
253+
"Change Target File",
254+
"Exit",
255+
]
256+
# the entry that is selected by default when the menu opens
257+
# this is overridden after the user picks an option
258+
# so that the default option is always the one that was last chosen
259+
default_response_option = (
260+
ResponseOptions.RECOMPILE_RUN
261+
if args.start
262+
else ResponseOptions.RECOMPILE_DOWNLOAD
263+
)
264+
265+
while True:
266+
try:
267+
if args.file is sys.stdin:
268+
await hub.race_disconnect(
247269
hub.race_power_button_press(
248-
questionary.select(
249-
f"Would you like to re-compile {os.path.basename(args.file.name)}?",
250-
response_options,
251-
default=(response_options[default_response_option]),
270+
questionary.press_any_key_to_continue(
271+
"The hub will stay connected and echo its output to the terminal. Press any key to exit."
252272
).ask_async()
253273
)
254274
)
275+
return
276+
response = await hub.race_disconnect(
277+
hub.race_power_button_press(
278+
questionary.select(
279+
f"Would you like to re-compile {os.path.basename(args.file.name)}?",
280+
response_options,
281+
default=(response_options[default_response_option]),
282+
).ask_async()
283+
)
284+
)
255285

256-
default_response_option = response_options.index(response)
286+
default_response_option = response_options.index(response)
257287

258-
match response_options.index(response):
288+
match response_options.index(response):
259289

260-
case ResponseOptions.RECOMPILE_RUN:
261-
with _get_script_path(args.file) as script_path:
262-
await hub.run(script_path, wait=True)
290+
case ResponseOptions.RECOMPILE_RUN:
291+
with _get_script_path(args.file) as script_path:
292+
await hub.run(script_path, wait=True)
263293

264-
case ResponseOptions.RECOMPILE_DOWNLOAD:
265-
with _get_script_path(args.file) as script_path:
266-
await hub.download(script_path)
294+
case ResponseOptions.RECOMPILE_DOWNLOAD:
295+
with _get_script_path(args.file) as script_path:
296+
await hub.download(script_path)
267297

268-
case ResponseOptions.RUN_STORED:
269-
if hub.fw_version < Version("3.2.0-beta.4"):
270-
print(
271-
"Running a stored program remotely is only supported in the hub firmware version >= v3.2.0."
272-
)
273-
else:
274-
await hub.start_user_program()
275-
await hub._wait_for_user_program_stop()
276-
277-
case ResponseOptions.CHANGE_TARGET_FILE:
278-
args.file.close()
279-
while True:
280-
try:
281-
args.file = open(
282-
await hub.race_disconnect(
283-
hub.race_power_button_press(
284-
questionary.path(
285-
"What file would you like to use?"
286-
).ask_async()
287-
)
298+
case ResponseOptions.RUN_STORED:
299+
if hub.fw_version < Version("3.2.0-beta.4"):
300+
print(
301+
"Running a stored program remotely is only supported in the hub firmware version >= v3.2.0."
302+
)
303+
else:
304+
await hub.start_user_program()
305+
await hub._wait_for_user_program_stop()
306+
307+
case ResponseOptions.CHANGE_TARGET_FILE:
308+
args.file.close()
309+
while True:
310+
try:
311+
args.file = open(
312+
await hub.race_disconnect(
313+
hub.race_power_button_press(
314+
questionary.path(
315+
"What file would you like to use?"
316+
).ask_async()
288317
)
289318
)
290-
break
291-
except FileNotFoundError:
292-
print("The file was not found. Please try again.")
293-
# send the new target file to the hub
294-
with _get_script_path(args.file) as script_path:
295-
await hub.download(script_path)
296-
297-
case _:
298-
return
299-
300-
except SyntaxError as e:
301-
print()
302-
print("A syntax error occurred while parsing your program:")
303-
print(e)
304-
print()
305-
306-
except HubPowerButtonPressedError:
307-
# This means the user pressed the button on the hub to re-start the
308-
# current program, so the menu was canceled and we are now printing
309-
# the hub stdout until the user program ends on the hub.
310-
try:
311-
await hub._wait_for_power_button_release()
312-
await hub._wait_for_user_program_stop()
313-
314-
except HubDisconnectError:
315-
hub = await reconnect_hub()
319+
)
320+
break
321+
except FileNotFoundError:
322+
print("The file was not found. Please try again.")
323+
# send the new target file to the hub
324+
with _get_script_path(args.file) as script_path:
325+
await hub.download(script_path)
326+
327+
case _:
328+
return
329+
330+
except SyntaxError as e:
331+
print()
332+
print("A syntax error occurred while parsing your program:")
333+
print(e)
334+
print()
335+
336+
except HubPowerButtonPressedError:
337+
# This means the user pressed the button on the hub to re-start the
338+
# current program, so the menu was canceled and we are now printing
339+
# the hub stdout until the user program ends on the hub.
340+
try:
341+
await hub._wait_for_power_button_release()
342+
await hub._wait_for_user_program_stop()
316343

317344
except HubDisconnectError:
318-
# let terminal cool off before making a new prompt
319-
await asyncio.sleep(0.3)
320345
hub = await reconnect_hub()
321346

347+
except HubDisconnectError:
348+
# let terminal cool off before making a new prompt
349+
await asyncio.sleep(0.3)
350+
hub = await reconnect_hub()
351+
352+
async def run(self, args: argparse.Namespace):
353+
322354
# Pick the right connection
323355
if args.conntype == "ble":
324356
from pybricksdev.ble import find_device as find_ble
@@ -382,13 +414,15 @@ def is_pybricks_usb(dev):
382414
await hub.download(script_path)
383415

384416
if args.stay_connected:
385-
await stay_connected_menu(hub)
417+
await self.stay_connected_menu(hub)
386418

387419
except SyntaxError as e:
388-
print("\nA syntax error occurred while parsing your program:")
389-
print(e, "\n")
420+
print()
421+
print("A syntax error occurred while parsing your program:")
422+
print(e)
423+
print()
390424
if args.stay_connected:
391-
await stay_connected_menu(hub)
425+
await self.stay_connected_menu(hub)
392426

393427
finally:
394428
await hub.disconnect()

0 commit comments

Comments
 (0)