From b1d7533d0705381e6ecee0fc95bb4c30f132a9fb Mon Sep 17 00:00:00 2001
From: Eric Cousineau <eric.cousineau@tri.global>
Date: Tue, 19 Nov 2019 14:58:37 -0500
Subject: [PATCH] py systems: Bind `System.SetDefaultContext` and
 `Simulator.reset_context`

---
 bindings/pydrake/systems/analysis_py.cc       |  7 +++++++
 .../pydrake/systems/framework_py_systems.cc   |  2 ++
 bindings/pydrake/systems/test/general_test.py | 19 +++++++++++++++++++
 3 files changed, 28 insertions(+)

diff --git a/bindings/pydrake/systems/analysis_py.cc b/bindings/pydrake/systems/analysis_py.cc
index b9bf10bcc3d6..c14045c566ea 100644
--- a/bindings/pydrake/systems/analysis_py.cc
+++ b/bindings/pydrake/systems/analysis_py.cc
@@ -105,6 +105,13 @@ PYBIND11_MODULE(analysis, m) {
             py_reference_internal, doc.Simulator.get_mutable_integrator.doc)
         .def("get_mutable_context", &Simulator<T>::get_mutable_context,
             py_reference_internal, doc.Simulator.get_mutable_context.doc)
+        .def("has_context", &Simulator<T>::has_context,
+            doc.Simulator.has_context.doc)
+        .def("reset_context", &Simulator<T>::reset_context, py::arg("context"),
+            // Keep alive, ownership: `context` keeps `self` alive.
+            py::keep_alive<2, 1>(), doc.Simulator.reset_context.doc)
+        // TODO(eric.cousineau): Bind `release_context` once some form of the
+        // PR RobotLocomotion/pybind11#33 lands. Presently, it fails.
         .def("reset_integrator",
             [](Simulator<T>* self,
                 std::unique_ptr<IntegratorBase<T>> integrator) {
diff --git a/bindings/pydrake/systems/framework_py_systems.cc b/bindings/pydrake/systems/framework_py_systems.cc
index 471e2a9bb52b..b2e226582d0f 100644
--- a/bindings/pydrake/systems/framework_py_systems.cc
+++ b/bindings/pydrake/systems/framework_py_systems.cc
@@ -328,6 +328,8 @@ struct Impl {
             doc.System.AllocateContext.doc)
         .def("CreateDefaultContext", &System<T>::CreateDefaultContext,
             doc.System.CreateDefaultContext.doc)
+        .def("SetDefaultContext", &System<T>::SetDefaultContext,
+            doc.System.SetDefaultContext.doc)
         .def("AllocateOutput",
             overload_cast_explicit<unique_ptr<SystemOutput<T>>>(
                 &System<T>::AllocateOutput),
diff --git a/bindings/pydrake/systems/test/general_test.py b/bindings/pydrake/systems/test/general_test.py
index 2ba8c6abf5ce..75dc3e9e2da6 100644
--- a/bindings/pydrake/systems/test/general_test.py
+++ b/bindings/pydrake/systems/test/general_test.py
@@ -111,6 +111,7 @@ def test_context_api(self):
             context.get_continuous_state_vector(), VectorBase)
         self.assertIsInstance(
             context.get_mutable_continuous_state_vector(), VectorBase)
+        system.SetDefaultContext(context)
 
         context = system.CreateDefaultContext()
         self.assertIsInstance(
@@ -436,6 +437,24 @@ def test_diagram_simulation(self):
                            xc_initial)
             self.assertTrue(np.allclose(xc, xc_expected))
 
+    def test_simulator_context_manipulation(self):
+        system = ConstantVectorSource([1])
+        # Use default-constructed context.
+        simulator = Simulator(system)
+        self.assertTrue(simulator.has_context())
+        context_default = simulator.get_mutable_context()
+        # WARNING: Once we call `simulator.reset_context()`, it will delete the
+        # context it currently owns, which is `context_default` in this case.
+        # BE CAREFUL IN SITUATIONS LIKE THIS!
+        # TODO(eric.cousineau): Bind `release_context()`, or migrate context
+        # usage to use `shared_ptr`.
+        context = system.CreateDefaultContext()
+        simulator.reset_context(context)
+        self.assertIs(context, simulator.get_mutable_context())
+        # WARNING: This will also invalidate `context`. Be careful!
+        simulator.reset_context(None)
+        self.assertFalse(simulator.has_context())
+
     def test_simulator_integrator_manipulation(self):
         system = ConstantVectorSource([1])