DISPATCH
cli.py
1 """Command line interface to f90nml.
2 
3 :copyright: Copyright 2017 Marshall Ward, see AUTHORS for details.
4 :license: Apache License, Version 2.0, see LICENSE for details
5 """
6 from __future__ import print_function
7 
8 import argparse
9 import json
10 import os
11 import sys
12 try:
13  from StringIO import StringIO # Python 2.x
14 except ImportError:
15  from io import StringIO # Python 3.x
16 try:
17  from collections import OrderedDict
18 except ImportError:
19  from ordereddict import OrderedDict
20 
21 import f90nml
22 try:
23  import yaml
24  has_yaml = True
25 
26  # Preserve ordering in YAML output
27  # https://stackoverflow.com/a/31609484/317172
28  represent_dict_order = (lambda self, data:
29  self.represent_mapping('tag:yaml.org,2002:map',
30  data.items()))
31  yaml.add_representer(OrderedDict, represent_dict_order)
32 
33 except ImportError:
34  has_yaml = False
35 
36 
37 def parse():
38  """Parse the command line input arguments."""
39  parser = argparse.ArgumentParser()
40 
41  parser.add_argument('--version', action='version',
42  version='f90nml {0}'.format(f90nml.__version__))
43 
44  parser.add_argument('--group', '-g', action='store')
45  parser.add_argument('--variable', '-v', action='append')
46  parser.add_argument('--patch', '-p', action='store_true')
47  parser.add_argument('--format', '-f', action='store')
48  parser.add_argument('--output', '-o', action='store')
49 
50  parser.add_argument('input', nargs='?')
51  parser.add_argument('output', nargs='?')
52 
53  if len(sys.argv) == 1:
54  parser.print_help()
55  sys.exit()
56 
57  args = parser.parse_args()
58 
59  input_fname = args.input
60  output_fname = args.output
61 
62  # Get input format
63  # TODO: Combine with output format
64  if input_fname:
65  _, input_ext = os.path.splitext(input_fname)
66  if input_ext == '.json':
67  input_fmt = 'json'
68  elif input_ext == '.yaml':
69  input_fmt = 'yaml'
70  else:
71  input_fmt = 'nml'
72  else:
73  input_fmt = 'nml'
74 
75  # Output format flag validation
76  valid_formats = ('json', 'yaml', 'nml')
77  if args.format and args.format not in valid_formats:
78  print('f90nml: error: format must be one of the following: {0}'
79  ''.format(valid_formats), file=sys.stderr)
80  sys.exit(-1)
81 
82  # Get output format
83  # TODO: Combine with input format
84  if not args.format:
85  if output_fname:
86  _, output_ext = os.path.splitext(output_fname)
87  if output_ext == '.json':
88  output_fmt = 'json'
89  elif output_ext in ('.yaml', '.yml'):
90  output_fmt = 'yaml'
91  else:
92  output_fmt = 'nml'
93  else:
94  output_fmt = 'nml'
95  else:
96  output_fmt = args.format
97 
98  # Confirm that YAML module is available
99  if (input_fmt == 'yaml' or output_fmt == 'yaml') and not has_yaml:
100  print('f90nml: error: YAML module could not be found.',
101  file=sys.stderr)
102  sys.exit(-1)
103 
104  # Do not patch non-namelist output
105  if any(fmt != 'nml' for fmt in (input_fmt, output_fmt)) and args.patch:
106  print('f90nml: error: Only namelist files can be patched.',
107  file=sys.stderr)
108  sys.exit(-1)
109 
110  # Read the input file
111  if input_fname:
112  if input_fmt in ('json', 'yaml'):
113  if input_fmt == 'json':
114  with open(input_fname) as input_file:
115  input_data = json.load(input_file)
116  elif input_ext == '.yaml':
117  with open(input_fname) as input_file:
118  input_data = yaml.safe_load(input_file)
119  else:
120  input_data = f90nml.read(input_fname)
121  else:
122  input_data = {}
123 
124  input_data = f90nml.Namelist(input_data)
125 
126  # Construct the update namelist
127  update_nml = {}
128  if args.variable:
129  if not args.group:
130  # Use the first available group
131  grp = list(input_data.keys())[0]
132  print('f90nml: warning: Assuming variables are in group \'{0}\'.'
133  ''.format(grp), file=sys.stderr)
134  else:
135  grp = args.group
136 
137  update_nml_str = '&{0} {1} /\n'.format(grp, ', '.join(args.variable))
138  update_io = StringIO(update_nml_str)
139  update_nml = f90nml.read(update_io)
140  update_io.close()
141 
142  # Target output
143  output_file = open(output_fname, 'w') if output_fname else sys.stdout
144 
145  if args.patch:
146  # We have to read the file twice for a patch. The main reason is
147  # to identify the default group, in case this is not provided.
148  # It could be avoided if a group is provided, but logically that could
149  # a mess that I do not want to sort out right now.
150  f90nml.patch(input_fname, update_nml, output_file)
151 
152  else:
153  # Update the input namelist directly
154  if update_nml:
155  try:
156  input_data[grp].update(update_nml[grp])
157  except KeyError:
158  input_data[grp] = update_nml[grp]
159 
160  # Write to output
161  if not args.patch:
162  if output_fmt in ('json', 'yaml'):
163  if output_fmt == 'json':
164  input_data = input_data.todict(complex_tuple=True)
165  json.dump(input_data, output_file,
166  indent=4, separators=(',', ': '))
167  output_file.write('\n')
168 
169  elif output_fmt == 'yaml':
170  input_data = input_data.todict(complex_tuple=True)
171  yaml.dump(input_data, output_file,
172  default_flow_style=False)
173  else:
174  # Default to namelist output
175  f90nml.write(input_data, output_file)
176 
177  # Cleanup
178  if output_file != sys.stdout:
179  output_file.close()
def parse()
Definition: cli.py:37
def patch(nml_path, nml_patch, out_path=None)
Definition: __init__.py:70
def read(nml_path)
Definition: __init__.py:12
def write(nml, nml_path, force=False, sort=False)
Definition: __init__.py:35