diff --git a/dragonflow/cli/df_model.py b/dragonflow/cli/df_model.py index baf75f10f..3b738ee5b 100644 --- a/dragonflow/cli/df_model.py +++ b/dragonflow/cli/df_model.py @@ -10,9 +10,14 @@ # License for the specific language governing permissions and limitations # under the License. +from __future__ import print_function + import abc +import argparse +import contextlib import re import six +import sys from jsonmodels import fields @@ -41,22 +46,31 @@ class ModelClass(object): @six.add_metaclass(abc.ABCMeta) class ModelPrinter(object): + def __init__(self, fh): + self._output = fh + @abc.abstractmethod def output_model(self, model_): pass class PlaintextPrinter(ModelPrinter): + def __init__(self, fh): + super(PlaintextPrinter, self).__init__(fh) + def output_model(self, model_): - print('\n-------------\n{}\n-------------'.format(model_.name)) + print('-------------', file=self._output) + print('{}'.format(model_.name), file=self._output) + print('-------------', file=self._output) for field in model_.fields: - if field.restrictions: - print('{name} - {type} {restriction}, {to_many}'.format( - name=field.name, type=field.type, - restriction=field.restrictions, to_many=field.to_many)) - else: - print('{name} - {type}, {to_many}'.format( - name=field.name, type=field.type, to_many=field.to_many)) + restriction_str = \ + ' {}'.format(field.restrictions) if field.restrictions else '' + print('{name} : {type}{restriction}, {to_many}'.format( + name=field.name, type=field.type, + restriction=restriction_str, + to_many="Many" if field.to_many else "One"), + file=self._output) + print('', file=self._output) class DfModelParser(object): @@ -65,8 +79,7 @@ class DfModelParser(object): return 'String', None elif isinstance(field, field_types.EnumField): field_type = type(field).__name__ - restrictions = [] - restrictions.extend(field._valid_values) + restrictions = list(field._valid_values) return field_type, restrictions elif isinstance(field, field_types.ReferenceField): model = field._model @@ -91,6 +104,8 @@ class DfModelParser(object): to_many = True field_type, restrictions = \ self._stringify_field_type(field.items_types[0]) + if isinstance(field, field_types.EnumListField): + restrictions = list(field._valid_values) else: to_many = False field_type, restrictions = self._stringify_field_type(field) @@ -105,7 +120,33 @@ class DfModelParser(object): printer.output_model(self._process_model(model)) -if __name__ == "__main__": - printer = PlaintextPrinter() +@contextlib.contextmanager +def smart_open(filename=None): + if filename and filename != '-': + fh = open(filename, 'w') + else: + fh = sys.stdout + + try: + yield fh + finally: + if fh is not sys.stdout: + fh.close() + + +def main(): + parser = argparse.ArgumentParser(description="Print Dragonflow schema") + group = parser.add_mutually_exclusive_group() + group.add_argument('--plaintext', help='Plaintext output (default)', + action='store_true') + parser.add_argument('-o', '--outfile', + help='Output to file (instead of stdout)') + args = parser.parse_args() parser = DfModelParser() - parser.parse_models(printer) + with smart_open(args.outfile) as fh: + printer = PlaintextPrinter(fh) + parser.parse_models(printer) + + +if __name__ == "__main__": + main()