Adds an API to retrieve keys. (#446)

This is useful for when we perform visualization. Signed-off-by: Arjo Chakravarty <[email protected]> Co-authored-by: Michael Carroll <[email protected]>
gazebosim · Oct 17, 2022 · efcf016 · efcf016
1 parent c29c69e
commit efcf016
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 1 deletion.
diff --git a/io/include/gz/common/DataFrame.hh b/io/include/gz/common/DataFrame.hh
@@ -68,6 +68,18 @@ namespace gz
         return this->storage.at(_key);
       }
 
+      /// \brief Retrieve all keys
+      /// \return A vector with keys
+      public: const std::vector<K> Keys() const
+      {
+        std::vector<K> keyList;
+        for (auto &[k, _]: this->storage)
+        {
+          keyList.push_back(k);
+        }
+        return keyList;
+      }
+
       /// \brief Data frame storage
       private: std::unordered_map<K, V> storage;
     };

diff --git a/io/src/DataFrame_TEST.cc b/io/src/DataFrame_TEST.cc
@@ -66,6 +66,9 @@ TEST(DataFrameTests, SimpleCSV)
       temperatureSession.value(), position);
   ASSERT_TRUE(temperature.has_value());
   EXPECT_DOUBLE_EQ(25.1, temperature.value());
+  auto keys = df.Keys();
+  ASSERT_EQ(keys.size(), 1);
+  ASSERT_EQ(keys[0], "temperature");
 }
 
 /////////////////////////////////////////////////
@@ -87,7 +90,8 @@ TEST(DataFrameTests, ComplexCSV)
   EXPECT_TRUE(df.Has("temperature"));
   EXPECT_TRUE(df.Has("humidity"));
   ASSERT_TRUE(df.Has("pressure"));
-
+  auto keys = df.Keys();
+  ASSERT_EQ(keys.size(), 3);
   const DataT &pressureData = df["pressure"];
   auto pressureSession = pressureData.CreateSession();
   const math::Vector3d position{36.80079505, -121.789472517, 0.8};