A cleaner interface for bindings.

2 files changed, 23 insertions, 21 deletions

diff --git a/traffic/pytraffic/bindings.cpp b/traffic/pytraffic/bindings.cpp
index 02c261358a..aacf206f95 100644
--- a/traffic/pytraffic/bindings.cpp
+++ b/traffic/pytraffic/bindings.cpp
@@ -15,6 +15,7 @@
 #include "std/string.hpp"
 #include "std/vector.hpp"
 
+#include "pyhelpers/vector_list_conversion.hpp"
 #include "pyhelpers/vector_uint8.hpp"
 
 #include <boost/python.hpp>
@@ -93,26 +94,39 @@ string RoadSegmentIdRepr(traffic::TrafficInfo::RoadSegmentId const & v)
   return ss.str();
 }
 
-vector<traffic::TrafficInfo::RoadSegmentId> GenerateTrafficKeys(string const & mwmPath)
+boost::python::list GenerateTrafficKeys(string const & mwmPath)
 {
   vector<traffic::TrafficInfo::RoadSegmentId> result;
   traffic::TrafficInfo::ExtractTrafficKeys(mwmPath, result);
-  return result;
+  return std_vector_to_python_list(result);
 }
 
-vector<uint8_t> GenerateTrafficValues(vector<traffic::TrafficInfo::RoadSegmentId> const & keys,
-                                      SegmentMapping const & segmentMapping)
+vector<uint8_t> GenerateTrafficValues(boost::python::list const & keys,
+                                      boost::python::dict const & segmentMappingDict)
 {
+  vector<traffic::TrafficInfo::RoadSegmentId> keysVec =
+      python_list_to_std_vector<traffic::TrafficInfo::RoadSegmentId>(keys);
+  SegmentMapping segmentMapping;
+
+  boost::python::list mappingKeys = segmentMappingDict.keys();
+  for (size_t i = 0; i < len(mappingKeys); ++i)
+  {
+    object curArg = segmentMappingDict[mappingKeys[i]];
+    if (curArg)
+      segmentMapping[extract<traffic::TrafficInfo::RoadSegmentId>(mappingKeys[i])] =
+          extract<SegmentSpeeds>(segmentMappingDict[mappingKeys[i]]);
+  }
+
   traffic::TrafficInfo::Coloring const knownColors = TransformToSpeedGroups(segmentMapping);
   traffic::TrafficInfo::Coloring coloring;
-  traffic::TrafficInfo::CombineColorings(keys, knownColors, coloring);
+  traffic::TrafficInfo::CombineColorings(keysVec, knownColors, coloring);
 
   vector<traffic::SpeedGroup> values(coloring.size());
 
   size_t i = 0;
   for (auto const & kv : coloring)
   {
-    ASSERT_EQUAL(kv.first, keys[i], ());
+    ASSERT_EQUAL(kv.first, keysVec[i], ());
     values[i] = kv.second;
     ++i;
   }
@@ -159,12 +173,6 @@ BOOST_PYTHON_MODULE(pytraffic)
       .value("Unknown", traffic::SpeedGroup::Unknown)
   ;
 
-  class_<traffic::TrafficInfo::Coloring>("Coloring")
-      .def(map_indexing_suite<traffic::TrafficInfo::Coloring>())
-  ;
-
-  class_<SegmentMapping>("SegmentMapping").def(map_indexing_suite<SegmentMapping>());
-
   def("load_classificator", LoadClassificator);
   def("generate_traffic_keys", GenerateTrafficKeys);
   def("generate_traffic_values", GenerateTrafficValues);
diff --git a/traffic/pytraffic/example.py b/traffic/pytraffic/example.py
index 6e7493cedb..79c2e8ff73 100644
--- a/traffic/pytraffic/example.py
+++ b/traffic/pytraffic/example.py
@@ -1,4 +1,4 @@
-from pytraffic import *
+from pytraffic import RoadSegmentId, SegmentSpeeds, load_classificator, generate_traffic_keys, generate_traffic_values
 import argparse
 
 parser = argparse.ArgumentParser(description='Example usage of pytraffic.')
@@ -12,14 +12,11 @@ if not options.path_to_classificator or not options.path_to_mwm:
 
 load_classificator(options.path_to_classificator)
 
-keys_list = [
+keys = [
   RoadSegmentId(0, 0, 0),
   RoadSegmentId(1, 0, 0),
   RoadSegmentId(1, 0, 1),
 ]
-keys_vec = RoadSegmentIdVec()
-for k in keys_list:
-  keys_vec.append(k)
 
 keys_from_mwm = generate_traffic_keys(options.path_to_mwm)
 
@@ -27,8 +24,5 @@ mapping = {
   RoadSegmentId(0, 0, 0):SegmentSpeeds(1.0, 2.0, 3.0),
   RoadSegmentId(1, 0, 1):SegmentSpeeds(4.0, 5.0, 6.0),
 }
-seg_map = SegmentMapping()
-for k, v in mapping.iteritems():
-  seg_map[k] = v
 
-buf = generate_traffic_values(keys_vec, seg_map)
+buf = generate_traffic_values(keys, mapping)