Streaming Key values with a custom formatter

gtsam/inference/Key.cpp

@@ -56,7 +56,7 @@ string _multirobotKeyFormatter(Key key) {
 
 /* ************************************************************************* */
 template<class CONTAINER>
-static void Print(const CONTAINER& keys, const string& s,
+void Print(const CONTAINER& keys, const string& s,
     const KeyFormatter& keyFormatter) {
   cout << s << " ";
   if (keys.empty())
@@ -83,6 +83,44 @@ void PrintKeySet(const KeySet& keys, const string& s,
     const KeyFormatter& keyFormatter) {
   Print(keys, s, keyFormatter);
 }
+
+/* ************************************************************************* */
+// Access to custom stream property.
+void *&key_formatter::property(ios_base &s) {
+  static int kUniqueIndex = ios_base::xalloc();
+  return s.pword(kUniqueIndex);
+}
+
+/* ************************************************************************* */
+// Store pointer to formatter in property.
+void key_formatter::set_property(ios_base &s, const KeyFormatter &f) {
+  property(s) = (void *)(&f);
+}
+
+/* ************************************************************************* */
+// Get pointer to formatter from property.
+KeyFormatter *key_formatter::get_property(ios_base &s) {
+  return (KeyFormatter *)(property(s));
+}
+
+/* ************************************************************************* */
+// Stream operator that will take a key_formatter and set the stream property.
+ostream &operator<<(ostream &os, const key_formatter &m) {
+  key_formatter::set_property(os, m.formatter_);
+  return os;
+}
+
+/* ************************************************************************* */
+// Stream operator that takes a StreamedKey and properly formats it
+ostream &operator<<(ostream &os, const StreamedKey &streamedKey) {
+  const KeyFormatter *formatter = key_formatter::get_property(os);
+  if (formatter == nullptr) {
+    formatter = &DefaultKeyFormatter;
+  }
+  os << (*formatter)(streamedKey.key_);
+  return (os);
+}
+
 /* ************************************************************************* */
 
 } // \namespace gtsam

gtsam/inference/Key.h

@@ -27,6 +27,8 @@
 
 #include <boost/function.hpp>
 
+#include <iosfwd>
+
 namespace gtsam {
 
 /// Typedef for a function to format a key, i.e. to convert it to a string
@@ -52,6 +54,34 @@ GTSAM_EXPORT std::string _multirobotKeyFormatter(gtsam::Key key);
 static const gtsam::KeyFormatter MultiRobotKeyFormatter =
     &_multirobotKeyFormatter;
 
+/// To use the key_formatter on Keys, they must be wrapped in a StreamedKey.
+struct StreamedKey {
+  const Key &key_;
+  explicit StreamedKey(const Key &key) : key_(key) {}
+  friend std::ostream &operator<<(std::ostream &, const StreamedKey &);
+};
+
+/**
+ * Output stream manipulator that will format gtsam::Keys according to the given
+ * KeyFormatter, as long as Key values are wrapped in a gtsam::StreamedKey.
+ * LabeledSymbol and Symbol values do not have to be wrapped.
+ * usage:
+ *   Key key = LabeledSymbol('x', 'A', 5); // cast to key type
+ *   cout << key_formatter(MultiRobotKeyFormatter) << StreamedKey(key);
+ */
+class key_formatter {
+ public:
+  explicit key_formatter(KeyFormatter v) : formatter_(v) {}
+  friend std::ostream &operator<<(std::ostream &, const key_formatter &);
+  friend std::ostream &operator<<(std::ostream &, const StreamedKey &);
+
+ private:
+  KeyFormatter formatter_;
+  static void *&property(std::ios_base &s);
+  static void set_property(std::ios_base &s, const KeyFormatter &f);
+  static KeyFormatter *get_property(std::ios_base &s);
+};
+
 /// Define collection type once and for all - also used in wrappers
 typedef FastVector<Key> KeyVector;
 

gtsam/inference/tests/testKey.cpp

@@ -22,6 +22,9 @@
 #include <CppUnitLite/TestHarness.h>
 
 #include <boost/assign/std/list.hpp> // for operator +=
+
+#include <sstream>
+
 using namespace boost::assign;
 using namespace std;
 using namespace gtsam;
@@ -41,17 +44,15 @@ TEST(Key, KeySymbolConversion) {
 template<int KeySize>
 Key KeyTestValue();
 
-template<>
-Key KeyTestValue<8>()
-{
+template <>
+Key KeyTestValue<8>() {
   return 0x6100000000000005;
-};
+}
 
-template<>
-Key KeyTestValue<4>()
-{
+template <>
+Key KeyTestValue<4>() {
   return 0x61000005;
-};
+}
 
 /* ************************************************************************* */
 TEST(Key, KeySymbolEncoding) {
@@ -68,12 +69,41 @@ TEST(Key, KeySymbolEncoding) {
 
 /* ************************************************************************* */
 TEST(Key, ChrTest) {
-  Key key = Symbol('c',3);
+  Symbol key('c', 3);
   EXPECT(Symbol::ChrTest('c')(key));
   EXPECT(!Symbol::ChrTest('d')(key));
 }
 
 /* ************************************************************************* */
-int main() { TestResult tr; return TestRegistry::runAllTests(tr); }
+// A custom (nonsensical) formatter.
+string myFormatter(Key key) {
+  return "special";
+}
+
+TEST(Key, Formatting) {
+  Symbol key('c', 3);
+  EXPECT("c3" == DefaultKeyFormatter(key));
+
+  // Try streaming keys, should be default-formatted.
+  stringstream ss;
+  ss << StreamedKey(key);
+  EXPECT("c3" == ss.str());
+
+  // use key_formatter with a function pointer
+  stringstream ss2;
+  ss2 << key_formatter(myFormatter) << StreamedKey(key);
+  EXPECT("special" == ss2.str());
+
+  // use key_formatter with a function object.
+  stringstream ss3;
+  ss3 << key_formatter(DefaultKeyFormatter) << StreamedKey(key);
+  EXPECT("c3" == ss3.str());
+}
+
+/* ************************************************************************* */
+int main() {
+  TestResult tr;
+  return TestRegistry::runAllTests(tr);
+}
 /* ************************************************************************* */