MeasurementService: Use a context manager to simplify ctrl-C handling (#110)

Author: bkeryan

examples/daqmx_analog_input_measurement/measurement.py

@@ -88,9 +88,8 @@ def main(verbose: int):
         level = logging.WARNING
     logging.basicConfig(format="%(asctime)s %(levelname)s: %(message)s", level=level)
 
-    daqmx_analog_input_measurement_service.host_service()
-    input("Press enter to close the measurement service.\n")
-    daqmx_analog_input_measurement_service.close_service()
+    with daqmx_analog_input_measurement_service.host_service():
+        input("Press enter to close the measurement service.\n")
 
 
 if __name__ == "__main__":

examples/dc_measurement_labviewUI/measurement.py

@@ -142,9 +142,8 @@ def main(verbose: int):
         level = logging.WARNING
     logging.basicConfig(format="%(asctime)s %(levelname)s: %(message)s", level=level)
 
-    dc_measurement_service.host_service()
-    input("Press enter to close the measurement service.\n")
-    dc_measurement_service.close_service()
+    with dc_measurement_service.host_service():
+        input("Press enter to close the measurement service.\n")
 
 
 if __name__ == "__main__":

examples/dc_measurement_measurementUI/measurement.py

@@ -142,9 +142,8 @@ def main(verbose: int):
         level = logging.WARNING
     logging.basicConfig(format="%(asctime)s %(levelname)s: %(message)s", level=level)
 
-    dc_measurement_service.host_service()
-    input("Press enter to close the measurement service.\n")
-    dc_measurement_service.close_service()
+    with dc_measurement_service.host_service():
+        input("Press enter to close the measurement service.\n")
 
 
 if __name__ == "__main__":

examples/sample_measurement/measurement.py

@@ -71,9 +71,8 @@ def main(verbose: int):
         level = logging.WARNING
     logging.basicConfig(format="%(asctime)s %(levelname)s: %(message)s", level=level)
 
-    sample_measurement_service.host_service()
-    input("Press enter to close the measurement service.\n")
-    sample_measurement_service.close_service()
+    with sample_measurement_service.host_service():
+        input("Press enter to close the measurement service.\n")
 
 
 if __name__ == "__main__":

ni_measurement_generator/ni_measurement_generator/templates/measurement.py.mako

@@ -48,9 +48,8 @@ def main(verbose: int):
         level = logging.WARNING
     logging.basicConfig(format="%(asctime)s %(levelname)s: %(message)s", level=level)
 
-    measurement_service.host_service()
-    input("Press enter to close the measurement service.\n")
-    measurement_service.close_service()
+    with measurement_service.host_service():
+        input("Press enter to close the measurement service.\n")
 
 
 if __name__ == "__main__":

ni_measurement_generator/tests/test_assets/example_renders/example.py

@@ -46,9 +46,8 @@ def main(verbose: int):
         level = logging.WARNING
     logging.basicConfig(format="%(asctime)s %(levelname)s: %(message)s", level=level)
 
-    measurement_service.host_service()
-    input("Press enter to close the measurement service.\n")
-    measurement_service.close_service()
+    with measurement_service.host_service():
+        input("Press enter to close the measurement service.\n")
 
 
 if __name__ == "__main__":

ni_measurement_service/_internal/service_manager.py

@@ -10,7 +10,6 @@
 from ni_measurement_service._internal.stubs.ni.measurements.measurementservice import (
     measurement_service_pb2_grpc,
 )
-from ni_measurement_service._internal.utilities import console_exit_functions
 from ni_measurement_service.measurement.info import MeasurementInfo, ServiceInfo
 
 
@@ -91,7 +90,6 @@ def start(
         _logger.info("Measurement service hosted on port: %s", port)
         self.discovery_client.register_measurement_service(port, service_info, measurement_info)
 
-        console_exit_functions.setup_unregister_on_console_close(self.stop)
         self.port = port
         return port
 

ni_measurement_service/_internal/utilities/console_exit_functions.py

@@ -1,5 +1,7 @@
 """Framework to host measurement service."""
 
+from __future__ import annotations
+
 from typing import Any, Callable
 
 from ni_measurement_service._internal import grpc_servicer
@@ -143,9 +145,14 @@ def _output(func):
 
         return _output
 
-    def host_service(self) -> None:
+    def host_service(self) -> MeasurementService:
         """Host the registered measurement method as gRPC measurement service.
 
+        Returns
+        -------
+            MeasurementService: Context manager that can be used with a with-statement to close
+            the service.
+
         Raises
         ------
             Exception: If register measurement methods not available.
@@ -160,8 +167,16 @@ def host_service(self) -> None:
             self.output_parameter_list,
             self.measure_function,
         )
-        return None
+        return self
 
     def close_service(self) -> None:
         """Close the Service after un-registering with discovery service and cleanups."""
         self.grpc_service.stop()
+
+    def __enter__(self) -> MeasurementService:
+        """Enter the runtime context related to the measurement service."""
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        """Exit the runtime context related to the measurement service."""
+        self.close_service()