pydrake: Add lcm.Subscriber implementation (RobotLocomotion#11175)

bindings/pydrake/BUILD.bazel

@@ -164,6 +164,7 @@ drake_pybind_library(
     py_deps = [
         ":module_py",
     ],
+    py_srcs = ["_lcm_extra.py"],
 )
 
 drake_pybind_library(

bindings/pydrake/_lcm_extra.py

@@ -0,0 +1,36 @@
+# See `ExecuteExtraPythonCode` in `pydrake_pybind.h` for usage details and
+# rationale.
+
+
+# This is a python reimplementation of drake::lcm::Subscriber.  We reimplement
+# (rather than bind) the C++ class because we want message() to be a python LCM
+# message, not a C++ object.
+class Subscriber:
+    """Subscribes to and stores a copy of the most recent message on a given
+    channel, for some message type.  This class does NOT provide any mutex
+    behavior for multi-threaded locking; it should only be used in cases where
+    the governing DrakeLcmInterface.HandleSubscriptions is called from the same
+    thread that owns all copies of this object.
+    """
+
+    def __init__(self, lcm, channel, lcm_type):
+        """Creates a subscriber and subscribes to `channel` on `lcm`.
+
+        Args:
+            lcm: LCM service instance (a pydrake.lcm.DrakeLcmInterface).
+            channel: LCM channel name.
+            lcm_type: Python class generated by lcmgen.
+        """
+        self.lcm_type = lcm_type
+        self.clear()
+        lcm.Subscribe(channel=channel, handler=self._handler)
+
+    def clear(self):
+        self.count = 0
+        self.raw = []
+        self.message = self.lcm_type()
+
+    def _handler(self, raw):
+        self.count += 1
+        self.raw = raw
+        self.message = self.lcm_type.decode(raw)

bindings/pydrake/lcm_py.cc

@@ -92,6 +92,8 @@ PYBIND11_MODULE(lcm, m) {
             doc.DrakeMockLcm.get_last_published_message.doc);
     // TODO(eric.cousineau): Add remaining methods.
   }
+
+  ExecuteExtraPythonCode(m);
 }
 
 }  // namespace pydrake

bindings/pydrake/systems/meshcat_visualizer.py

@@ -15,7 +15,7 @@
 from drake import lcmt_viewer_load_robot
 from pydrake.common.eigen_geometry import Quaternion, Isometry3
 from pydrake.geometry import DispatchLoadMessage, SceneGraph
-from pydrake.lcm import DrakeMockLcm
+from pydrake.lcm import DrakeMockLcm, Subscriber
 from pydrake.math import RigidTransform, RotationMatrix
 from pydrake.systems.framework import (
     AbstractValue, LeafSystem, PublishEvent, TriggerType
@@ -250,16 +250,16 @@ def load(self):
         self.vis[self.prefix].delete()
 
         # Intercept load message via mock LCM.
-        def handle_load_robot(raw):
-            load_robot_msgs.append(lcmt_viewer_load_robot.decode(raw))
-        load_robot_msgs = []
         mock_lcm = DrakeMockLcm()
-        mock_lcm.Subscribe(
+        mock_lcm_subscriber = Subscriber(
+            lcm=mock_lcm,
             channel="DRAKE_VIEWER_LOAD_ROBOT",
-            handler=handle_load_robot)
+            lcm_type=lcmt_viewer_load_robot)
         DispatchLoadMessage(self._scene_graph, mock_lcm)
         mock_lcm.HandleSubscriptions(0)
-        load_robot_msg = load_robot_msgs[0]
+        assert mock_lcm_subscriber.count > 0
+        load_robot_msg = mock_lcm_subscriber.message
+
         # Translate elements to `meshcat`.
         for i in range(load_robot_msg.num_links):
             link = load_robot_msg.link[i]

bindings/pydrake/systems/test/lcm_test.py

@@ -10,7 +10,7 @@
 
 from robotlocomotion import header_t, quaternion_t
 
-from pydrake.lcm import DrakeLcm, DrakeMockLcm
+from pydrake.lcm import DrakeLcm, DrakeMockLcm, Subscriber
 from pydrake.systems.analysis import Simulator
 from pydrake.systems.framework import (
     AbstractValue, BasicVector, DiagramBuilder, LeafSystem)
@@ -145,23 +145,23 @@ def test_publisher(self):
         dut = mut.LcmPublisherSystem.Make(
             channel="TEST_CHANNEL", lcm_type=quaternion_t, lcm=lcm,
             publish_period=0.1)
+        subscriber = Subscriber(lcm, "TEST_CHANNEL", quaternion_t)
         model_message = self._model_message()
         self._fix_and_publish(dut, AbstractValue.Make(model_message))
-        raw = lcm.get_last_published_message("TEST_CHANNEL")
-        actual_message = quaternion_t.decode(raw)
-        self.assert_lcm_equal(actual_message, model_message)
+        lcm.HandleSubscriptions(0)
+        self.assert_lcm_equal(subscriber.message, model_message)
 
     def test_publisher_cpp(self):
         lcm = DrakeMockLcm()
         dut = mut.LcmPublisherSystem.Make(
             channel="TEST_CHANNEL", lcm_type=quaternion_t, lcm=lcm,
             use_cpp_serializer=True)
+        subscriber = Subscriber(lcm, "TEST_CHANNEL", quaternion_t)
         model_message = self._model_message()
         model_value = self._model_value_cpp()
         self._fix_and_publish(dut, model_value)
-        raw = lcm.get_last_published_message("TEST_CHANNEL")
-        actual_message = quaternion_t.decode(raw)
-        self.assert_lcm_equal(actual_message, model_message)
+        lcm.HandleSubscriptions(0)
+        self.assert_lcm_equal(subscriber.message, model_message)
 
     def test_connect_lcm_scope(self):
         builder = DiagramBuilder()

bindings/pydrake/test/lcm_test.py

@@ -1,6 +1,6 @@
 import unittest
 
-from pydrake.lcm import DrakeLcm, DrakeLcmInterface, DrakeMockLcm
+from pydrake.lcm import DrakeLcm, DrakeLcmInterface, DrakeMockLcm, Subscriber
 
 from robotlocomotion import quaternion_t
 
@@ -38,6 +38,24 @@ def test_mock_lcm_roundtrip(self):
         dut.HandleSubscriptions(0)
         self.assertEqual(self.count, 1)
 
+    def test_subscriber(self):
+        mock = DrakeMockLcm()
+        dut = Subscriber(lcm=mock, channel="CHANNEL", lcm_type=quaternion_t)
+        mock.Publish(channel="CHANNEL", buffer=self.quat.encode())
+        self.assertEqual(dut.count, 0)
+        self.assertEqual(len(dut.raw), 0)
+        self.assertEqual(dut.message.w, 0)
+        self.assertEqual(dut.message.x, 0)
+        self.assertEqual(dut.message.y, 0)
+        self.assertEqual(dut.message.z, 0)
+        mock.HandleSubscriptions(0)
+        self.assertEqual(dut.count, 1)
+        self.assertEqual(dut.raw, self.quat.encode())
+        self.assertEqual(dut.message.w, self.quat.w)
+        self.assertEqual(dut.message.x, self.quat.x)
+        self.assertEqual(dut.message.y, self.quat.y)
+        self.assertEqual(dut.message.z, self.quat.z)
+
     def test_mock_lcm_get_last_published_message(self):
         dut = DrakeMockLcm()
         dut.Publish(channel="TEST_CHANNEL", buffer=self.quat.encode())

lcm/drake_lcm_interface.h

@@ -225,16 +225,16 @@ class Subscriber final {
   DRAKE_DEFAULT_COPY_AND_MOVE_AND_ASSIGN(Subscriber)
 
   /**
-   * Subscribes to the (non-empty) @p channel_name on the given (non-null)
+   * Subscribes to the (non-empty) @p channel on the given (non-null)
    * @p lcm instance.  The `lcm` pointer is only used during construction; it
    * is not retained by this object.  When a undecodable message is received,
    * @p on_error handler is invoked; when `on_error` is not provided, an
    * exception will be thrown instead.
    */
-  Subscriber(DrakeLcmInterface* lcm, const std::string& channel_name,
+  Subscriber(DrakeLcmInterface* lcm, const std::string& channel,
              std::function<void()> on_error = {}) {
     subscription_ = drake::lcm::Subscribe<Message>(
-        lcm, channel_name, [data = data_](const Message& message) {
+        lcm, channel, [data = data_](const Message& message) {
           data->message = message;
           data->count++;
         }, std::move(on_error));