az resource-graph --include displayNames fails on ubuntu

src/azure/cli/_argparse.py

@@ -124,7 +124,7 @@ def add_command(self,
             else:
                 v = aliases.pop().strip('<> ')
             target, _ = _read_arg(aliases[0])
-            kw.update({_read_arg(a)[0]: (target, v, req) for a in aliases})
+            kw.update({_read_arg(a)[0]: (target, v, req, aliases) for a in aliases})
             ad.append(('/'.join(aliases), desc, req))
 
     #pylint: disable=too-many-branches
@@ -217,7 +217,7 @@ def not_global(a):
                 parsed.positional.append(n)
             n = next_n
 
-        required_args = [x for x, _, req in expected_kwargs.values() if req]
+        required_args = [x for x, _, req, _ in expected_kwargs.values() if req]
         for a in required_args:
             try:
                 parsed[a]
@@ -279,24 +279,35 @@ def _display_usage(self, nouns, noun_map, out=sys.stdout):
             logger.debug('Expected documentation at %s', doc_file)
 
     def _display_completions(self, noun_map, arguments, out=sys.stdout):
+
         for a in self.complete_args:
             arguments.remove(a)
 
         command_candidates = set([k for k in noun_map if not k.startswith('$')])
-        if command_candidates and not arguments[-1].startswith('-'):
-            command_candidates = set([c for c in command_candidates if c.startswith(arguments[-1])])
+
+        last_arg = arguments[-1]
+        if command_candidates and not last_arg.startswith('-'):
+            command_candidates = set([c for c in command_candidates if c.startswith(last_arg)])
 
         kwargs = noun_map.get('$kwargs') or []
-        args_candidates = set(('--' if len(a) > 1 else '-') + a for a in kwargs)
-        if arguments[-1].startswith('-'):
+        args_candidates = set()
+        arguments_set = set(arguments)
+
+        #handle a messy part, that if a short name is used, then the long name should
+        #be excluded. Say, for '-a/--arg', if '-a' is used in command, '--arg' should
+        #not in the candidate list.
+        for a in kwargs:
+            alias = kwargs[a][3]
+            if not [x for x in alias if x in arguments_set]:
+                args_candidates.update(set(alias))
+
+        if last_arg.startswith('-'):
             # TODO: We don't have enough metadata about the command to do parameter value
             # completion (yet). This should only apply to value arguments, not flag arguments
-            if arguments[-1] in args_candidates:
+            if last_arg in args_candidates:
                 args_candidates = set()
             else:
-                args_candidates = set([c for c in args_candidates if c.startswith(arguments[-1])])
-        else:
-            args_candidates = args_candidates.difference(arguments)
+                args_candidates = set([c for c in args_candidates if c.startswith(last_arg)])
 
         candidates = command_candidates.union(args_candidates)
 

src/azure/cli/tests/test_argparse.py

@@ -113,32 +113,29 @@ def test_required_args(self):
 
         self.assertIsNone(p.execute('n1 -b x'.split()).result)
 
-    def test_args_completion(self):
+    def test_specify_output_format(self):
         p = ArgumentParser('test')
         p.add_command(lambda a, b: (a, b), 'n1', args=[('--arg -a', '', True), ('-b <v>', '', False)])
-        io = StringIO()
 
-        p.execute('n1 - --complete'.split(), 
-                               show_usage=False,
-                               show_completions=True,
-                               out=io)
-        candidates = util.normalize_newlines(io.getvalue())
-        self.assertEqual(candidates, '--arg\n-a\n-b\n')
+        cmd_res = p.execute('n1 -a x'.split())
+        self.assertEqual(cmd_res.output_format, None)
 
-        io = StringIO()
-        p.execute('n1 --a --complete'.split(), 
-                               show_usage=False,
-                               out=io)
-        candidates = util.normalize_newlines(io.getvalue())
-        self.assertEqual(candidates, '--arg\n')
+        cmd_res = p.execute('n1 -a x --output json'.split())
+        self.assertEqual(cmd_res.output_format, 'json')
 
-        io = StringIO()
-        p.execute('n --a --complete'.split(), 
-                               show_usage=False,
-                               show_completions=True,
-                               out=io)
-        candidates = util.normalize_newlines(io.getvalue())
-        self.assertEqual(candidates, 'n1\n')
+        cmd_res = p.execute('n1 -a x --output table'.split())
+        self.assertEqual(cmd_res.output_format, 'table')
+
+        cmd_res = p.execute('n1 -a x --output text'.split())
+        self.assertEqual(cmd_res.output_format, 'text')
+
+        # Invalid format
+        cmd_res = p.execute('n1 -a x --output unknown'.split())
+        self.assertEqual(cmd_res.output_format, None)
+
+        # Invalid format
+        cmd_res = p.execute('n1 -a x --output'.split())
+        self.assertEqual(cmd_res.output_format, None)
 
     def test_specify_output_format(self):
         p = ArgumentParser('test')
@@ -164,5 +161,69 @@ def test_specify_output_format(self):
         cmd_result = p.execute('n1 -a x --output'.split())
         self.assertEqual(cmd_result.output_format, None)
 
+    def test_args_completion(self):
+        p = ArgumentParser('test')
+        p.add_command(lambda a, b: (a, b), 'n1', args=[('--arg -a', '', True), ('-b <v>', '', False)])
+
+        # Can't use "with StringIO() as ...", as Python2/StringIO doesn't have __exit__.
+        io = StringIO()
+        p.execute('n1 - --complete'.split(), 
+                               show_usage=False,
+                               show_completions=True,
+                               out=io)
+        candidates = util.normalize_newlines(io.getvalue())
+        io.close()
+        self.assertEqual(candidates, '--arg\n-a\n-b\n')
+
+        #matching '--arg for '--a'
+        io=StringIO()
+        p.execute('n1 --a --complete'.split(), 
+                               show_usage=False,
+                               out=io)
+        candidates = util.normalize_newlines(io.getvalue())
+        io.close()
+        self.assertEqual(candidates, '--arg\n')
+
+        #matching 'n1' for 'n'
+        io = StringIO()
+        p.execute('n --complete'.split(), 
+                               show_usage=False,
+                               show_completions=True,
+                               out=io)
+        candidates = util.normalize_newlines(io.getvalue())
+        io.close()
+        self.assertEqual(candidates, 'n1\n')
+
+        #if --arg is used, then both '-a' and "--arg" should not be in the 
+        #candidate list 
+        io = StringIO()
+        p.execute('n1 --arg hello - --complete'.split(), 
+                               show_usage=False,
+                               show_completions=True,
+                               out=io)
+        candidates = util.normalize_newlines(io.getvalue())
+        io.close()
+        self.assertEqual(candidates, '-b\n')
+
+        #if all argument are used, candidate list is empty 
+        io = StringIO()
+        p.execute('n1 -a -b --complete'.split(), 
+                               show_usage=False,
+                               show_completions=True,
+                               out=io)
+        candidates = util.normalize_newlines(io.getvalue())
+        io.close()
+        self.assertEqual(candidates, '\n')
+
+        #if at parameter value level, get nothing for N.Y.I.
+        io = StringIO()
+        p.execute('n1 -a --complete'.split(), 
+                               show_usage=False,
+                               show_completions=True,
+                               out=io)
+        candidates = util.normalize_newlines(io.getvalue())
+        io.close()
+        self.assertEqual(candidates, '\n')
+
 if __name__ == '__main__':
     unittest.main()

src/azure/cli/tests/test_output.py

@@ -3,11 +3,12 @@
 import unittest
 
 try:
+    # on Python2, we like to use the module "StringIO" rather "io" so to 
+    # avoid "print" errors like: TypeError: string argument expected, got 'str'
+    from StringIO import StringIO
+except ImportError:
     # Python 3
     from io import StringIO
-except ImportError:
-    # Python 2
-    from StringIO import StringIO
 
 from azure.cli._output import (OutputProducer, OutputFormatException, format_json, format_table, format_list, format_text,
                                 ListOutput)