Incorporate PR feedback

* Add empty record type into conversion script
autorope · Dec 23, 2020 · 3b4e42a · 3b4e42a
1 parent 35d6ac8
commit 3b4e42a
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/donkeycar/parts/tub_v2.py b/donkeycar/parts/tub_v2.py
@@ -13,8 +13,6 @@ class Tub(object):
     A datastore to store sensor data in a key, value format. \n
     Accepts str, int, float, image_array, image, and array data types.
     """
-    # static variable for empty records
-    empty_record = {'__empty__': True}
 
     def __init__(self, base_path, inputs=[], types=[], metadata=[],
                  max_catalog_len=1000, read_only=False):
@@ -35,8 +33,6 @@ def write_record(self, record=None):
         """
         Can handle various data types including images.
         """
-        if not record:
-            record = self.empty_record
         contents = dict()
         for key, value in record.items():
             if value is None:

diff --git a/scripts/convert_to_tub_v2.py b/scripts/convert_to_tub_v2.py
@@ -22,6 +22,8 @@
 
 
 def convert_to_tub_v2(paths, output_path):
+    empty_record = {'__empty__': True}
+
     if type(paths) is str:
         paths = [paths]
     legacy_tubs = [LegacyTub(path) for path in paths]
@@ -30,8 +32,11 @@ def convert_to_tub_v2(paths, output_path):
 
     for legacy_tub in legacy_tubs:
         if not output_tub:
-            output_tub = Tub(output_path, legacy_tub.inputs,
-                             legacy_tub.types, list(legacy_tub.meta.items()))
+            # add input and type for empty records recording
+            inputs = legacy_tub.inputs + ['__empty__']
+            types = legacy_tub.types + ['boolean']
+            output_tub = Tub(output_path, inputs, types,
+                             list(legacy_tub.meta.items()))
 
         record_paths = legacy_tub.gather_records()
         bar = IncrementalBar('Converting', max=len(record_paths))
@@ -58,7 +63,7 @@ def convert_to_tub_v2(paths, output_path):
                     # until the next valid record.
                     while previous_index < current_index:
                         idx = output_tub.manifest.current_index
-                        output_tub.write_record()
+                        output_tub.write_record(empty_record)
                         output_tub.delete_record(idx)
                         previous_index += 1
                 bar.next()