Add a model format updater

scripts/tests/test_translate_model.py

@@ -24,9 +24,35 @@
 from scripts import translate_model  # type: ignore # noqa (module hack)
 
 
-class TestTranslate(unittest.TestCase):
+class TestNormalize(unittest.TestCase):
 
-  def test_icu(self) -> None:
+  def test_old_format_input(self) -> None:
+    model = {'a:x': 48, 'a:y': 21, 'b:x': 2, 'b:z': 89}
+    expect = {'a': {'x': 48, 'y': 21}, 'b': {'x': 2, 'z': 89}}
+    result = translate_model.normalize(model)
+    self.assertDictEqual(result, expect)
+
+  def test_new_format_input(self) -> None:
+    model = {'a': {'x': 48, 'y': 21}, 'b': {'x': 2, 'z': 89}}
+    result = translate_model.normalize(model)
+    self.assertDictEqual(result, model)
+
+  def test_broken_input1(self) -> None:
+    model = {'a:x': 23, 'b': {'x': 37, 'y': 18}}
+    with self.assertRaises(Exception) as cm:
+      translate_model.normalize(model)
+    self.assertTrue('Unsupported model format' in str(cm.exception))
+
+  def test_broken_input2(self) -> None:
+    model = {'b': {'x': 37, 'y': {'z': 123}}}
+    with self.assertRaises(Exception) as cm:
+      translate_model.normalize(model)
+    self.assertTrue('Unsupported model format' in str(cm.exception))
+
+
+class TestTranslateICU(unittest.TestCase):
+
+  def test_standard(self) -> None:
     model = {'a': {'x': 12, 'y': 88}, 'b': {'x': 47, 'z': 13}}
     expect = '''
 jaml {

scripts/translate_model.py

@@ -11,9 +11,20 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Translates a model JSON file to another format, such as ICU Resource Bundle."""
+"""Translates a model JSON file to another format, such as ICU Resource Bundle.
+
+Example usage:
+
+$ python translate_model.py --format=icu model.json > icurb.txt
+
+You can also use this script to update the model files older than v0.5.0 to make
+it work with the latest version.
+
+$ python translate_model.py --format=json old-model.json > new-model.json
+"""
 
 import argparse
+import itertools
 import json
 import typing
 
@@ -46,15 +57,61 @@ def translate_icu(model: typing.Dict[str, typing.Dict[str, int]]) -> str:
   return output
 
 
+def normalize(
+    model: typing.Dict[str,
+                       typing.Any]) -> typing.Dict[str, typing.Dict[str, int]]:
+  """Updates a model to the latest format. Does nothing if it's updated already.
+
+  Args:
+    model: A model.
+  Returns:
+    An updated model.
+  """
+  is_old_format = all([isinstance(v, int) for v in model.values()])
+  if is_old_format:
+    output = {}
+    sorted_items = sorted(model.items(), key=lambda x: x[0])
+    groups = itertools.groupby(sorted_items, key=lambda x: x[0].split(':')[0])
+    for group in groups:
+      output[group[0]] = dict(
+          (item[0].split(':')[-1], item[1]) for item in group[1])
+    return output
+  try:
+    assert (all([
+        isinstance(v, int)
+        for groups in model.values()
+        for v in groups.values()
+    ])), 'Scores should be integers'
+  except (AssertionError, AttributeError) as e:
+    raise Exception('Unsupported model format:', e)
+  else:
+    return model
+
+
 def main() -> None:
-  parser = argparse.ArgumentParser(description=__doc__)
+  DEFAULT_FORMAT = 'json'
+  parser = argparse.ArgumentParser(
+      description=__doc__, formatter_class=argparse.RawTextHelpFormatter)
   parser.add_argument(
       'model', help='File path for the JSON format model file.', type=str)
+  parser.add_argument(
+      '--format',
+      help=f'Target format (default: {DEFAULT_FORMAT})',
+      type=str,
+      default=DEFAULT_FORMAT,
+      choices={DEFAULT_FORMAT, 'icu'})
   args = parser.parse_args()
   model_path: str = args.model
+  format: str = args.format
   with open(model_path) as f:
     model = json.load(f)
-  print(translate_icu(model))
+  model = normalize(model)
+  if format == 'json':
+    print(json.dumps(model, ensure_ascii=False, separators=(',', ':')))
+  elif format == 'icu':
+    print(translate_icu(model))
+  else:
+    pass
 
 
 if __name__ == '__main__':