Support for a flag to inject environment variables

fire/core.py

@@ -321,6 +321,7 @@ def _Fire(component, args, context, name=None):
   show_completion = parsed_flag_args.completion
   show_help = parsed_flag_args.help
   show_trace = parsed_flag_args.trace
+  use_environment = parsed_flag_args.use_environment
 
   # component can be a module, class, routine, object, etc.
   if component is None:
@@ -363,7 +364,7 @@ def _Fire(component, args, context, name=None):
         filename, lineno = inspectutils.GetFileAndLine(component)
 
         component, consumed_args, remaining_args, capacity = _CallCallable(
-            component, remaining_args)
+            component, remaining_args, use_environment)
 
         # Update the trace.
         if isclass:
@@ -524,7 +525,7 @@ def _GetMember(component, args):
   raise FireError('Could not consume arg:', arg)
 
 
-def _CallCallable(fn, args):
+def _CallCallable(fn, args, use_environment=None):
   """Calls the function fn by consuming args from args.
 
   Args:
@@ -536,14 +537,14 @@ def _CallCallable(fn, args):
     remaining_args: The remaining args that haven't been consumed yet.
     capacity: Whether the call could have taken additional args.
   """
-  parse = _MakeParseFn(fn)
+  parse = _MakeParseFn(fn, use_environment=use_environment)
   (varargs, kwargs), consumed_args, remaining_args, capacity = parse(args)
 
   result = fn(*varargs, **kwargs)
   return result, consumed_args, remaining_args, capacity
 
 
-def _MakeParseFn(fn):
+def _MakeParseFn(fn, use_environment=None):
   """Creates a parse function for fn.
 
   Args:
@@ -567,6 +568,16 @@ def _ParseFn(args):
     """Parses the list of `args` into (varargs, kwargs), remaining_args."""
     kwargs, remaining_args = _ParseKeywordArgs(args, all_args, fn_spec.varkw)
 
+    # Enrich with environment variables
+    if use_environment is not None:
+      for arg in fn_spec.args[len(remaining_args):]:
+        env_var_name = use_environment + arg
+        env_var_name_upper = use_environment + arg.upper()
+        env_value = os.environ.get(env_var_name,
+                                   os.environ.get(env_var_name_upper))
+        if arg not in kwargs and env_value:
+          kwargs[arg] = env_value
+
     # Note: _ParseArgs modifies kwargs.
     parsed_args, kwargs, remaining_args, capacity = _ParseArgs(
         fn_spec.args, fn_spec.defaults, num_required_args, kwargs,

fire/core_test.py

@@ -18,6 +18,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import os
+
 from fire import core
 from fire import test_components as tc
 from fire import testutils
@@ -72,6 +74,24 @@ def testInteractiveModeVariablesWithName(self, mock_embed):
     self.assertEqual(variables['D'], tc.WithDefaults)
     self.assertIsInstance(variables['trace'], trace.FireTrace)
 
+  def testUsingEnvironment(self):
+    os.environ['count'] = '4'
+    res = core.Fire(tc.WithDefaults,
+                    command=['double', '--', '--use_environment'])
+    assert res == 8
+
+  def testUsingEnvironmentDefaultsToUpper(self):
+    os.environ['COUNT'] = '4'
+    res = core.Fire(tc.WithDefaults,
+                    command=['double', '--', '--use_environment'])
+    assert res == 8
+
+  def testUsingEnvironmentNamespaced(self):
+    os.environ['TEST_NS_count'] = '5'
+    res = core.Fire(tc.WithDefaults,
+                    command=['double', '--', '--use_environment', 'TEST_NS_'])
+    assert res == 10
+
   def testImproperUseOfHelp(self):
     # This should produce a warning explaining the proper use of help.
     with self.assertRaisesFireExit(2, 'The proper way to show help.*Usage:'):

fire/parser.py

@@ -23,13 +23,20 @@
 
 
 def CreateParser():
+  """Creates the parser for the flags for Python fire
+
+  Returns:
+    An instance of argparse.ArgumentParser with the flags already added.
+
+  """
   parser = argparse.ArgumentParser(add_help=False)
   parser.add_argument('--verbose', '-v', action='store_true')
   parser.add_argument('--interactive', '-i', action='store_true')
   parser.add_argument('--separator', default='-')
   parser.add_argument('--completion', action='store_true')
   parser.add_argument('--help', '-h', action='store_true')
   parser.add_argument('--trace', '-t', action='store_true')
+  parser.add_argument('--use_environment', nargs='?', default=None, const='')
   # TODO: Consider allowing name to be passed as an argument.
   return parser