From 7dec8f2355bfa179cb6a0c212263a781ffcd7d97 Mon Sep 17 00:00:00 2001
From: Felipe Martins Diel <41558831+felipediel@users.noreply.github.com>
Date: Tue, 12 May 2020 16:19:49 -0300
Subject: [PATCH] Handle ReadError in the CLI (#365)

* Handle ReadError

* Make it more pythonic

* Clean up
---
 cli/broadlink_cli | 94 ++++++++++++++++++++++++-----------------------
 1 file changed, 48 insertions(+), 46 deletions(-)

diff --git a/cli/broadlink_cli b/cli/broadlink_cli
index 5045c5c..ca3f23e 100755
--- a/cli/broadlink_cli
+++ b/cli/broadlink_cli
@@ -6,8 +6,10 @@ import codecs
 import time
 
 import broadlink
+from broadlink.exceptions import ReadError
 
 TICK = 32.84
+TIMEOUT = 30
 IR_TOKEN = 0x26
 
 
@@ -85,7 +87,7 @@ parser.add_argument("--learnfile", help="save learned command to a specified fil
 parser.add_argument("--durations", action="store_true",
                     help="use durations in micro seconds instead of the Broadlink format")
 parser.add_argument("--convert", action="store_true", help="convert input data to durations")
-parser.add_argument("--joinwifi", nargs=2, help="Args are SSID PASSPHRASE to configure Broadlink device with");
+parser.add_argument("--joinwifi", nargs=2, help="Args are SSID PASSPHRASE to configure Broadlink device with")
 parser.add_argument("data", nargs='*', help="Data to send or convert")
 args = parser.parse_args()
 
@@ -128,27 +130,29 @@ if args.send:
     dev.send_data(data)
 if args.learn or args.learnfile:
     dev.enter_learning()
-    data = None
     print("Learning...")
-    timeout = 30
-    while (data is None) and (timeout > 0):
-        time.sleep(2)
-        timeout -= 2
-        data = dev.check_data()
-    if data:
-        learned = format_durations(to_microseconds(bytearray(data))) \
-            if args.durations \
-            else ''.join(format(x, '02x') for x in bytearray(data))
-        if args.learn:
-            print(learned)
-            decode_hex = codecs.getdecoder("hex_codec")
-            print("Base64: " + str(base64.b64encode(decode_hex(learned)[0])))
-        if args.learnfile:
-            print("Saving to {}".format(args.learnfile))
-            with open(args.learnfile, "w") as text_file:
-                text_file.write(learned)
+    for second in range(TIMEOUT):
+        try:
+            data = dev.check_data()
+        except ReadError:
+            time.sleep(1)
+        else:
+            break
     else:
         print("No data received...")
+        exit(1)
+
+    learned = format_durations(to_microseconds(bytearray(data))) \
+        if args.durations \
+        else ''.join(format(x, '02x') for x in bytearray(data))
+    if args.learn:
+        print(learned)
+        decode_hex = codecs.getdecoder("hex_codec")
+        print("Base64: " + str(base64.b64encode(decode_hex(learned)[0])))
+    if args.learnfile:
+        print("Saving to {}".format(args.learnfile))
+        with open(args.learnfile, "w") as text_file:
+            text_file.write(learned)
 if args.check:
     if dev.check_power():
         print('* ON *')
@@ -194,13 +198,11 @@ if args.rfscanlearn:
     dev.sweep_frequency()
     print("Learning RF Frequency, press and hold the button to learn...")
 
-    timeout = 20
-
-    while (not dev.check_frequency()) and (timeout > 0):
+    for second in range(TIMEOUT):
         time.sleep(1)
-        timeout -= 1
-
-    if timeout <= 0:
+        if dev.check_frequency():
+            break
+    else:
         print("RF Frequency not found")
         dev.cancel_sweep_frequency()
         exit(1)
@@ -214,26 +216,26 @@ if args.rfscanlearn:
 
     dev.find_rf_packet()
 
-    data = None
-    timeout = 20
-
-    while (data is None) and (timeout > 0):
-        time.sleep(1)
-        timeout -= 1
-        data = dev.check_data()
-
-    if data:
-        print("Found RF Frequency - 2 of 2!")
-        learned = format_durations(to_microseconds(bytearray(data))) \
-            if args.durations \
-            else ''.join(format(x, '02x') for x in bytearray(data))
-        if args.learnfile is None:
-            print(learned)
-            decode_hex = codecs.getdecoder("hex_codec")
-            print("Base64: {}".format(str(base64.b64encode(decode_hex(learned)[0]))))
-        if args.learnfile is not None:
-            print("Saving to {}".format(args.learnfile))
-            with open(args.learnfile, "w") as text_file:
-                text_file.write(learned)
+    for second in range(TIMEOUT):
+        try:
+            data = dev.check_data()
+        except ReadError:
+            time.sleep(1)
+        else:
+            break
     else:
         print("No data received...")
+        exit(1)
+
+    print("Found RF Frequency - 2 of 2!")
+    learned = format_durations(to_microseconds(bytearray(data))) \
+        if args.durations \
+        else ''.join(format(x, '02x') for x in bytearray(data))
+    if args.learnfile is None:
+        print(learned)
+        decode_hex = codecs.getdecoder("hex_codec")
+        print("Base64: {}".format(str(base64.b64encode(decode_hex(learned)[0]))))
+    if args.learnfile is not None:
+        print("Saving to {}".format(args.learnfile))
+        with open(args.learnfile, "w") as text_file:
+            text_file.write(learned)