[traffic] Python bindings for binary keys section.

2 files changed, 43 insertions, 15 deletions

diff --git a/traffic/pytraffic/bindings.cpp b/traffic/pytraffic/bindings.cpp
index e44256b2d2..8408055859 100644
--- a/traffic/pytraffic/bindings.cpp
+++ b/traffic/pytraffic/bindings.cpp
@@ -101,13 +101,10 @@ boost::python::list GenerateTrafficKeys(string const & mwmPath)
   return std_vector_to_python_list(result);
 }
 
-vector<uint8_t> GenerateTrafficValues(boost::python::list const & keys,
+vector<uint8_t> GenerateTrafficValues(vector<traffic::TrafficInfo::RoadSegmentId> 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)
   {
@@ -119,14 +116,14 @@ vector<uint8_t> GenerateTrafficValues(boost::python::list const & keys,
 
   traffic::TrafficInfo::Coloring const knownColors = TransformToSpeedGroups(segmentMapping);
   traffic::TrafficInfo::Coloring coloring;
-  traffic::TrafficInfo::CombineColorings(keysVec, knownColors, coloring);
+  traffic::TrafficInfo::CombineColorings(keys, knownColors, coloring);
 
   vector<traffic::SpeedGroup> values(coloring.size());
 
   size_t i = 0;
   for (auto const & kv : coloring)
   {
-    ASSERT_EQUAL(kv.first, keysVec[i], ());
+    ASSERT_EQUAL(kv.first, keys[i], ());
     values[i] = kv.second;
     ++i;
   }
@@ -137,6 +134,24 @@ vector<uint8_t> GenerateTrafficValues(boost::python::list const & keys,
   return buf;
 }
 
+vector<uint8_t> GenerateTrafficValuesFromList(boost::python::list const & keys,
+                                              boost::python::dict const & segmentMappingDict)
+{
+  vector<traffic::TrafficInfo::RoadSegmentId> keysVec =
+      python_list_to_std_vector<traffic::TrafficInfo::RoadSegmentId>(keys);
+
+  return GenerateTrafficValues(keysVec, segmentMappingDict);
+}
+
+vector<uint8_t> GenerateTrafficValuesFromBinary(vector<uint8_t> const & keysBlob,
+                                                boost::python::dict const & segmentMappingDict)
+{
+  vector<traffic::TrafficInfo::RoadSegmentId> keys;
+  traffic::TrafficInfo::DeserializeTrafficKeys(keysBlob, keys);
+
+  return GenerateTrafficValues(keys, segmentMappingDict);
+}
+
 void LoadClassificator(string const & classifPath)
 {
   GetPlatform().SetResourceDir(classifPath);
@@ -182,5 +197,6 @@ BOOST_PYTHON_MODULE(pytraffic)
 
   def("load_classificator", LoadClassificator);
   def("generate_traffic_keys", GenerateTrafficKeys);
-  def("generate_traffic_values", GenerateTrafficValues);
+  def("generate_traffic_values_from_list", GenerateTrafficValuesFromList);
+  def("generate_traffic_values_from_binary", GenerateTrafficValuesFromBinary);
 }
diff --git a/traffic/pytraffic/example.py b/traffic/pytraffic/example.py
index 87a8ff1d88..385215bb7f 100644
--- a/traffic/pytraffic/example.py
+++ b/traffic/pytraffic/example.py
@@ -1,12 +1,21 @@
-from pytraffic import RoadSegmentId, SegmentSpeeds, load_classificator, generate_traffic_keys, generate_traffic_values
+from __future__ import print_function
+from pytraffic import (RoadSegmentId,
+                       SegmentSpeeds,
+                       load_classificator,
+                       generate_traffic_keys,
+                       generate_traffic_values_from_binary,
+                       generate_traffic_values_from_list)
 import argparse
 
 parser = argparse.ArgumentParser(description='Example usage of pytraffic.')
-parser.add_argument("--path_to_classificator", dest="path_to_classificator", help="Path to the directory that contains classificator.txt.")
+parser.add_argument("--path_to_classificator", dest="path_to_classificator",
+                    help="Path to the directory that contains classificator.txt.")
 parser.add_argument("--path_to_mwm", dest="path_to_mwm", help="Path to the target mwm file.")
+parser.add_argument("--path_to_keys", dest="path_to_keys",
+                    help="Path to serialized key data (i.e. the \"traffic\" section).")
 
 options = parser.parse_args()
-if not options.path_to_classificator or not options.path_to_mwm:
+if not options.path_to_classificator or not options.path_to_mwm or not options.path_to_keys:
   parser.print_help()
   exit()
 
@@ -19,17 +28,20 @@ keys = [
 ]
 
 fid, idx, dir = keys[2].fid, keys[2].idx, keys[2].dir
-print fid, idx, dir
+print(fid, idx, dir)
 
 keys_from_mwm = generate_traffic_keys(options.path_to_mwm)
 
 seg_speeds = SegmentSpeeds(1.0, 2.0, 3.0)
 ws, wrs, w = seg_speeds.weighted_speed, seg_speeds.weighted_ref_speed, seg_speeds.weight
-print ws, wrs, w
+print(ws, wrs, w)
 
 mapping = {
-  RoadSegmentId(0, 0, 0):SegmentSpeeds(1.0, 2.0, 3.0),
-  RoadSegmentId(1, 0, 1):SegmentSpeeds(4.0, 5.0, 6.0),
+  RoadSegmentId(0, 0, 0): SegmentSpeeds(1.0, 2.0, 3.0),
+  RoadSegmentId(1, 0, 1): SegmentSpeeds(4.0, 5.0, 6.0),
 }
 
-buf = generate_traffic_values(keys, mapping)
+buf1 = generate_traffic_values_from_list(keys, mapping)
+
+with open(options.path_to_keys, "rb") as bin_data:
+  buf2 = generate_traffic_values_from_binary(bin_data.read(), {})