From ffb2f8aa96544a795d1a52bc8a9aec85e5d3be8e Mon Sep 17 00:00:00 2001
From: Laurens Valk <laurens@pybricks.com>
Date: Fri, 7 Jan 2022 13:37:03 +0100
Subject: [PATCH] pybricks.pupdevices.UltrasonicSensor: Move cleanup example.

Also generalize the example so that it can be used on all hubs.
---
 doc/main/hubs/primehub.rst                |  7 -----
 doc/main/pupdevices/ultrasonicsensor.rst  |  6 +++++
 examples/pup/hub_primehub/cleanup.py      | 33 -----------------------
 examples/pup/sensor_ultrasonic/cleanup.py | 25 +++++++++++++++++
 4 files changed, 31 insertions(+), 40 deletions(-)
 delete mode 100644 examples/pup/hub_primehub/cleanup.py
 create mode 100644 examples/pup/sensor_ultrasonic/cleanup.py

diff --git a/doc/main/hubs/primehub.rst b/doc/main/hubs/primehub.rst
index 6cc4aec..a194296 100644
--- a/doc/main/hubs/primehub.rst
+++ b/doc/main/hubs/primehub.rst
@@ -235,10 +235,3 @@ Turning the hub off
 
 .. literalinclude::
     ../../../examples/pup/hub_primehub/system_shutdown.py
-
-Clean up after program stop
-*****************************************
-
-.. literalinclude::
-    ../../../examples/pup/hub_primehub/cleanup.py
-
diff --git a/doc/main/pupdevices/ultrasonicsensor.rst b/doc/main/pupdevices/ultrasonicsensor.rst
index 2b6858d..16ff578 100644
--- a/doc/main/pupdevices/ultrasonicsensor.rst
+++ b/doc/main/pupdevices/ultrasonicsensor.rst
@@ -36,3 +36,9 @@ Gradually change the brightness of the lights
 
 .. literalinclude::
     ../../../examples/pup/sensor_ultrasonic/math.py
+
+Turn off the lights when the program ends
+**********************************************
+
+.. literalinclude::
+    ../../../examples/pup/sensor_ultrasonic/cleanup.py
diff --git a/examples/pup/hub_primehub/cleanup.py b/examples/pup/hub_primehub/cleanup.py
deleted file mode 100644
index e7de877..0000000
--- a/examples/pup/hub_primehub/cleanup.py
+++ /dev/null
@@ -1,33 +0,0 @@
-from pybricks.hubs import PrimeHub
-from pybricks.parameters import Color, Port
-from pybricks.pupdevices import ColorSensor, UltrasonicSensor
-
-# Initialize the hub
-hub = PrimeHub()
-
-# Initialize the sensors
-ultrasonic_sensor = UltrasonicSensor(Port.E)
-color_sensor = ColorSensor(Port.F)
-
-# Turn on upper halves
-ultrasonic_sensor.lights.on([50, 50, 0, 0])
-
-
-# Your robot's code goes here
-def main():
-    while True:
-        color = color_sensor.color()
-        if color == Color.RED:
-            hub.display.text("R")
-        else:
-            hub.display.text("?")
-
-
-# Wrap the code in try-finally to ensure the cleanup code is always executed
-# no matter if the program gets stopped gracefully by the user or by an error
-try:
-    main()
-finally:
-    # The cleanup code goes here
-    color_sensor.lights.off()
-    ultrasonic_sensor.lights.off()
diff --git a/examples/pup/sensor_ultrasonic/cleanup.py b/examples/pup/sensor_ultrasonic/cleanup.py
new file mode 100644
index 0000000..30d1504
--- /dev/null
+++ b/examples/pup/sensor_ultrasonic/cleanup.py
@@ -0,0 +1,25 @@
+from pybricks.parameters import Port
+from pybricks.pupdevices import UltrasonicSensor
+from pybricks.tools import wait
+
+# Initialize the sensor.
+sensor = UltrasonicSensor(Port.A)
+
+# Turn on the upper lights.
+sensor.lights.on([50, 50, 0, 0])
+
+
+def main():
+    while True:
+        print("Main program is running.")
+        wait(500)
+
+
+# Wrap the main code in try/finally so that the cleanup code always runs
+# when the program ends, even if an exception was raised.
+try:
+    main()
+finally:
+    # The cleanup code goes here.
+    print("Cleaning up.")
+    sensor.lights.off()