cpp_types: Add ability to map C++ types to Python to enable simple te…

…mplating.

bindings/pydrake/BUILD.bazel

@@ -49,7 +49,7 @@ drake_pybind_library(
     cc_srcs = ["common_py.cc"],
     package_info = PACKAGE_INFO,
     py_deps = [
-        ":util_py",
+        "//bindings/pydrake/util:module_shim_py",
     ],
     py_srcs = [
         "__init__.py",

bindings/pydrake/util/BUILD.bazel

@@ -6,9 +6,12 @@ load(
     "//tools:drake.bzl",
     "drake_cc_googletest",
     "drake_cc_library",
+    "drake_cc_binary",
 )
 load(
     "//tools/skylark:pybind.bzl",
+    "drake_pybind_cc_googletest",
+    "drake_pybind_library",
     "get_drake_pybind_installs",
     "get_pybind_package_info",
 )
@@ -52,7 +55,9 @@ drake_py_library(
     imports = PACKAGE_INFO.py_imports,
 )
 
-PYBIND_LIBRARIES = []
+PYBIND_LIBRARIES = [
+    ":cpp_types_py",
+]
 
 PY_LIBRARIES = [
     ":module_py",
@@ -66,6 +71,30 @@ drake_py_library(
     deps = PYBIND_LIBRARIES + PY_LIBRARIES,
 )
 
+# ODR does not matter, because the singleton will be stored in Python.
+drake_cc_library(
+    name = "cpp_types_pybind",
+    srcs = ["cpp_types_pybind.cc"],
+    hdrs = ["cpp_types_pybind.h"],
+    deps = [
+        ":type_pack",
+        "@pybind11",
+    ],
+)
+
+drake_pybind_library(
+    name = "cpp_types_py",
+    cc_deps = [
+        ":cpp_types_pybind",
+    ],
+    cc_srcs = ["cpp_types_py.cc"],
+    package_info = PACKAGE_INFO,
+    py_deps = [
+        ":module_py",
+    ],
+    py_srcs = ["cpp_types.py"],
+)
+
 install(
     name = "install",
     targets = PY_LIBRARIES,
@@ -100,4 +129,10 @@ drake_py_test(
     ],
 )
 
+drake_pybind_cc_googletest(
+    name = "cpp_types_test",
+    cc_deps = [":cpp_types_pybind"],
+    py_deps = [":cpp_types_py"],
+)
+
 add_lint_tests()

bindings/pydrake/util/cpp_types.py

@@ -0,0 +1,55 @@
+from __future__ import absolute_import, print_function
+
+import ctypes
+import numpy as np
+
+"""
+@file
+Defines a mapping between Python and C++ types, and provides canonical Python
+types as they relate to C++.
+"""
+
+# Define these first, as they are used in `cpp_types_py.cc`
+# (transitively, `cpp_types_pybind.cc`).
+
+
+def _get_type_name(t):
+    # Gets type name as a string.
+    prefix = t.__module__ + "."
+    if prefix == "__builtin__.":
+        prefix = ""
+    return prefix + t.__name__
+
+
+class _StrictMap(object):
+    def __init__(self):
+        self._values = dict()
+
+    def _strict_key(self, key):
+        # Ensures keys are strictly scoped to the values (for literals).
+        return (type(key), key)
+
+    def add(self, key, value):
+        skey = self._strict_key(key)
+        assert skey not in self._values, "Already added: {}".format(skey)
+        self._values[skey] = value
+
+    def get(self, key, default):
+        skey = self._strict_key(key)
+        return self._values.get(skey, default)
+
+
+# Load and import type registry.
+from ._cpp_types_py import _type_registry  # noqa
+
+
+def get_types_canonical(param):
+    """Gets the canonical types for a set of Python types (canonical as in
+    how they relate to C++ types. """
+    return tuple(map(_type_registry.GetPyTypeCanonical, param))
+
+
+def get_type_names(param):
+    """Gets the canonical type names for a set of Python types (canonical as in
+    how they relate to C++ types. """
+    return tuple(map(_type_registry.GetName, param))

bindings/pydrake/util/cpp_types_py.cc

@@ -0,0 +1,15 @@
+#include <pybind11/pybind11.h>
+
+#include "drake/bindings/pydrake/util/cpp_types_pybind.h"
+
+using drake::pydrake::internal::TypeRegistry;
+
+PYBIND11_MODULE(_cpp_types_py, m) {
+  py::class_<TypeRegistry> type_registry_cls(m, "_TypeRegistry");
+  type_registry_cls
+    .def(py::init<>())
+    .def("GetPyTypeCanonical", &TypeRegistry::GetPyTypeCanonical)
+    .def("GetName", &TypeRegistry::GetName);
+  // Create singleton instance.
+  m.attr("_type_registry") = type_registry_cls();
+}

bindings/pydrake/util/cpp_types_pybind.cc

@@ -0,0 +1,171 @@
+#include "drake/bindings/pydrake/util/cpp_types_pybind.h"
+
+#include <pybind11/eval.h>
+
+const char kModule[] = "pydrake.util.cpp_types";
+
+namespace drake {
+namespace pydrake {
+
+namespace internal {
+
+TypeRegistry::TypeRegistry() {
+  // Import modules into `locals_`.
+  globals_ = py::module::import(kModule).attr("__dict__");
+  py_to_py_canonical_ = eval("_StrictMap")();
+
+  RegisterCommon();
+  RegisterLiterals();
+}
+
+const TypeRegistry& TypeRegistry::GetPyInstance() {
+  auto m = py::module::import(kModule);
+  py::object type_registry_py = m.attr("_type_registry");
+  const TypeRegistry* type_registry =
+      py::cast<const TypeRegistry*>(type_registry_py);
+  return *type_registry;
+}
+
+py::object TypeRegistry::DoGetPyType(const std::type_info& tinfo) const {
+  // Check if it's a custom-registered type.
+  size_t cpp_key = std::type_index(tinfo).hash_code();
+  auto iter = cpp_to_py_.find(cpp_key);
+  if (iter != cpp_to_py_.end()) {
+    return iter->second;
+  } else {
+    // Get from pybind11-registered types.
+    // WARNING: Internal API :(
+    auto* info = py::detail::get_type_info(tinfo);
+    if (!info) {
+      throw std::runtime_error("Unknown type!");
+    }
+    return py::reinterpret_borrow<py::object>(
+        py::handle(reinterpret_cast<PyObject*>(info->type)));
+  }
+}
+
+py::object TypeRegistry::GetPyTypeCanonical(py::object py_type) const {
+  // Get registered canonical type if there is a mapping; otherwise return
+  // original type.
+  return py_to_py_canonical_.attr("get")(py_type, py_type);
+}
+
+py::str TypeRegistry::GetName(py::object py_type) const {
+  py::object py_type_canonical = GetPyTypeCanonical(py_type);
+  py::object name = py_name_.attr("get")(py_type_canonical);
+  // Assume this is a Python type.
+  if (name.is(py::none())) {
+    name = eval("_get_type_name")(py_type_canonical);
+  }
+  return name;
+}
+
+py::object TypeRegistry::eval(const std::string& expr) const {
+  return py::eval(expr, globals_, locals_);
+}
+
+void TypeRegistry::Register(
+      const std::vector<size_t>& cpp_keys,
+      py::tuple py_types, const std::string& name) {
+  py::object py_canonical = py_types[0];
+  for (size_t cpp_key : cpp_keys) {
+    assert(cpp_to_py_.find(cpp_key) == cpp_to_py_.end());
+    cpp_to_py_[cpp_key] = py_canonical;
+  }
+  for (auto py_type : py_types) {
+    py_to_py_canonical_.attr("add")(py_type, py_canonical);
+  }
+  py_name_[py_canonical] = name;
+}
+
+template <typename T>
+void TypeRegistry::RegisterType(
+    py::tuple py_types) {
+  py::str name = eval("_get_type_name")(py_types[0]);
+  Register({type_hash<T>()}, py_types, name.cast<std::string>());
+}
+
+void TypeRegistry::RegisterCommon() {
+  // Make mappings for C++ RTTI to Python types.
+  // Unfortunately, this is hard to obtain from `pybind11`.
+  RegisterType<bool>(eval("bool,"));
+  RegisterType<std::string>(eval("str,"));
+  RegisterType<double>(eval("float, np.double, ctypes.c_double"));
+  RegisterType<float>(eval("np.float32, ctypes.c_float"));
+  RegisterType<int>(eval("int, np.int32, ctypes.c_int32"));
+  RegisterType<uint32_t>(eval("np.uint32, ctypes.c_uint32"));
+  RegisterType<int64_t>(eval("np.int64, ctypes.c_int64"));
+  // For supporting generic Python types.
+  RegisterType<py::object>(eval("object,"));
+}
+
+class TypeRegistry::LiteralHelper {
+ public:
+  explicit LiteralHelper(TypeRegistry* type_registry)
+    : type_registry_(type_registry) {}
+
+  void RegisterLiterals() {
+    RegisterSequence(Render(std::integer_sequence<bool, false, true>{}));
+    // Register `int` (and `uint` as an alias for positive values).
+    constexpr int i_max = 100;
+    RegisterSequence(MakeSequence<int, -i_max, -1>());
+    RegisterSequence(
+        MakeSequence<int, 0, i_max>(),
+        {MakeSequence<uint, 0, i_max, int>()});
+  }
+
+ private:
+  template <typename T>
+  struct Sequence {
+    std::vector<size_t> keys;
+    std::vector<T> values;
+  };
+
+  template <typename T, typename IsAliasFor = T, T... Values>
+  Sequence<IsAliasFor> Render(std::integer_sequence<T, Values...>) {
+    return Sequence<IsAliasFor>{
+      {type_hash<std::integral_constant<T, Values>>()...},
+      {Values...}};
+  }
+
+  template <typename T, T Start, T End, typename IsAliasFor = T>
+  Sequence<IsAliasFor> MakeSequence() {
+    constexpr T Count = End - Start + 1;
+    auto seq = sequence_transform(
+        constant_add<T, Start>{}, std::make_integer_sequence<T, Count>{});
+    return Render<T, IsAliasFor>(seq);
+  }
+
+  template <typename T>
+  void RegisterSequence(
+      const Sequence<T>& seq,
+      std::vector<Sequence<T>> alias_set = {}) {
+    for (int i = 0; i < seq.keys.size(); ++i) {
+      // Get alias types.
+      std::vector<size_t> cpp_keys{seq.keys[i]};
+      for (const auto& alias : alias_set) {
+        assert(seq.values[i] == alias.values[i]);
+        cpp_keys.push_back(alias.keys[i]);
+      }
+      // Register.
+      T value = seq.values[i];
+      py::object py_value = py::cast(value);
+      type_registry_->Register(
+          cpp_keys,
+          py::make_tuple(py_value),
+          py::str(py_value).cast<std::string>());
+    }
+  }
+
+  TypeRegistry* type_registry_;
+};
+
+void TypeRegistry::RegisterLiterals() {
+  // Register a subset of literals.
+  LiteralHelper(this).RegisterLiterals();
+}
+
+}  // namespace internal
+
+}  // namespace pydrake
+}  // namespace drake