From c9cb95ea53244f7c3503c38099c1e8caa428e4cc Mon Sep 17 00:00:00 2001
From: Pete Gadomski <pete.gadomski@gmail.com>
Date: Tue, 9 Sep 2014 21:09:23 +0000
Subject: [PATCH] Programmable filter can add dimensions to context

Using one or more `add_dimension` options on a programmable filter will
instruct the filter to add those dimensions to the primary
PointContext. This way, the programmable filter can set dimensions that
might be not present on the PointContext at first.
---
 doc/stages/filters.programmable.rst        |  4 ++
 include/pdal/filters/Programmable.hpp      |  2 +
 src/filters/Programmable.cpp               | 15 +++++++
 test/unit/plang/ProgrammableFilterTest.cpp | 46 ++++++++++++++++++++++
 4 files changed, 67 insertions(+)

diff --git a/doc/stages/filters.programmable.rst b/doc/stages/filters.programmable.rst
index ecf6798efa..69b2b0fb2c 100644
--- a/doc/stages/filters.programmable.rst
+++ b/doc/stages/filters.programmable.rst
@@ -21,6 +21,8 @@ This example scales the Z value by a factor of ten (of course, you could also us
 
 Note that the function always returns `True`. If the function returned `False`, an error would be thrown and the translation shut down.
 
+If you want to write a dimension that might not be available, use can use one or more `add_dimension` options.
+
 To filter points based on a `Python`_ function, use the :ref:`filters.predicate` filter.
 
 Example
@@ -75,6 +77,8 @@ function
 source
   The literal `Python`_ code to execute, when the script option is not being used.
 
+add_dimension
+  The name of a dimension to add to the pipeline that does not already exist.
 
 .. _Python: http://python.org/
 .. _NumPy: http://www.numpy.org/
diff --git a/include/pdal/filters/Programmable.hpp b/include/pdal/filters/Programmable.hpp
index af991292f7..bc3f57e6af 100644
--- a/include/pdal/filters/Programmable.hpp
+++ b/include/pdal/filters/Programmable.hpp
@@ -69,8 +69,10 @@ class PDAL_DLL Programmable : public Filter
     std::string m_source;
     std::string m_module;
     std::string m_function;
+    Dimension::IdList m_addDimensions;
 
     virtual void processOptions(const Options& options);
+    virtual void addDimensions(PointContext ctx);
     virtual void ready(PointContext ctx);
     virtual void filter(PointBuffer& buf);
     virtual void done(PointContext ctx);
diff --git a/src/filters/Programmable.cpp b/src/filters/Programmable.cpp
index 1baf8d5b97..196c3fd208 100644
--- a/src/filters/Programmable.cpp
+++ b/src/filters/Programmable.cpp
@@ -63,6 +63,21 @@ void Programmable::processOptions(const Options& options)
             options.getValueOrThrow<std::string>("filename"));
     m_module = options.getValueOrThrow<std::string>("module");
     m_function = options.getValueOrThrow<std::string>("function");
+
+    auto addDims = options.getOptions("add_dimension");
+    for (auto it = addDims.cbegin(); it != addDims.cend(); ++it)
+    {
+        m_addDimensions.push_back(Dimension::id(it->getValue<std::string>()));
+    }
+}
+
+
+void Programmable::addDimensions(PointContext ctx)
+{
+    for (auto it = m_addDimensions.cbegin(); it != m_addDimensions.cend(); ++it)
+    {
+        ctx.registerDim(*it);
+    }
 }
 
 
diff --git a/test/unit/plang/ProgrammableFilterTest.cpp b/test/unit/plang/ProgrammableFilterTest.cpp
index 56963bb026..6805b38767 100644
--- a/test/unit/plang/ProgrammableFilterTest.cpp
+++ b/test/unit/plang/ProgrammableFilterTest.cpp
@@ -130,6 +130,52 @@ BOOST_AUTO_TEST_CASE(pipeline)
     }
 }
 
+
+BOOST_AUTO_TEST_CASE(add_dimension)
+{
+    Bounds<double> bounds(0.0, 0.0, 0.0, 1.0, 1.0, 1.0);
+
+    Options ops;
+    ops.add("bounds", bounds);
+    ops.add("num_points", 10);
+    ops.add("mode", "ramp");
+
+    drivers::faux::Reader reader(ops);
+
+    Option source("source", "import numpy\n"
+        "def myfunc(ins,outs):\n"
+        "  outs['Intensity'] = np.zeros(ins['X'].size, dtype=numpy.uint16) + 1\n"
+        "  outs['PointSourceId'] = np.zeros(ins['X'].size, dtype=numpy.uint16) + 2\n"
+        "  return True\n"
+    );
+    Option module("module", "MyModule");
+    Option function("function", "myfunc");
+    Option intensity("add_dimension", "Intensity");
+    Option scanDirection("add_dimension", "PointSourceId");
+    Options opts;
+    opts.add(source);
+    opts.add(module);
+    opts.add(function);
+    opts.add(intensity);
+    opts.add(scanDirection);
+
+    filters::Programmable filter(opts);
+    filter.setInput(&reader);
+
+    PointContext ctx;
+    filter.prepare(ctx);
+    PointBufferSet pbSet = filter.execute(ctx);
+    BOOST_CHECK_EQUAL(pbSet.size(), 1);
+    PointBufferPtr buf = *pbSet.begin();
+
+    for (unsigned int i = 0; i < buf->size(); ++i)
+    {
+        BOOST_CHECK_EQUAL(buf->getFieldAs<uint16_t>(Dimension::Id::Intensity, i), 1);
+        BOOST_CHECK_EQUAL(buf->getFieldAs<uint16_t>(Dimension::Id::PointSourceId, i), 2);
+    }
+}
+
+
 BOOST_AUTO_TEST_SUITE_END()
 
 #endif  // PDAL_HAVE_PYTHON