Provide completion scripts for fire.Fire(fn).

dbieber · dbieber · commit 9eb788624958 · 2017-03-31T17:42:49.000Z
Copybara generated commit for Python Fire. PiperOrigin-RevId: 151823094 Change-Id: I03073378734c83f27be69c9ea6fb237bb4f1e08e Reviewed-on: https://team-review.git.corp.google.com/66328 Reviewed-by: David Bieber <dbieber@google.com>
diff --git a/fire/completion.py b/fire/completion.py
@@ -210,6 +210,12 @@ def _Commands(component, depth=3):
     Tuples, each tuple representing one possible command for this CLI.
     Only traverses the member DAG up to a depth of depth.
   """
+  if inspect.isroutine(component) or inspect.isclass(component):
+    for completion in Completions(component):
+      yield (completion,)
+  if inspect.isroutine(component):
+    return  # Don't descend into routines.
+
   if depth < 1:
     return
 
@@ -219,10 +225,5 @@ def _Commands(component, depth=3):
 
     yield (member_name,)
 
-    if inspect.isroutine(member) or inspect.isclass(member):
-      for completion in Completions(member):
-        yield (member_name, completion)
-      continue  # Don't descend into routines.
-
     for command in _Commands(member, depth - 1):
       yield (member_name,) + command
diff --git a/fire/completion_test.py b/fire/completion_test.py
@@ -98,6 +98,21 @@ def testDeepDictScript(self):
     self.assertIn('level3', script)
     self.assertNotIn('level4', script)  # The default depth is 3.
 
+  def testFnScript(self):
+    script = completion.Script('identity', tc.identity)
+    self.assertIn('--arg1', script)
+    self.assertIn('--arg2', script)
+    self.assertIn('--arg3', script)
+    self.assertIn('--arg4', script)
+
+  def testClassScript(self):
+    script = completion.Script('', tc.MixedDefaults)
+    self.assertIn('ten', script)
+    self.assertIn('sum', script)
+    self.assertIn('identity', script)
+    self.assertIn('--alpha', script)
+    self.assertIn('--beta', script)
+
   def testNonStringDictCompletions(self):
     completions = completion.Completions({
         10: 'green',
diff --git a/fire/test_components.py b/fire/test_components.py
@@ -21,7 +21,7 @@
 import six
 
 if six.PY3:
-  from fire import test_components_py3 as py3  # pylint: disable=unused-import,no-name-in-module
+  from fire import test_components_py3 as py3  # pylint: disable=unused-import,no-name-in-module,g-import-not-at-top
 
 
 def identity(arg1, arg2, arg3=10, arg4=20, *arg5, **arg6):
