Initial draft of azdev extension update-index

src/azdev/commands.py

@@ -51,6 +51,7 @@ def operation_group(name):
         g.command('add', 'add_extension')
         g.command('remove', 'remove_extension')
         g.command('list', 'list_extensions')
+        g.command('update-index', 'update_extension_index')
 
     # TODO: implement
     # with CommandGroup(self, operation_group('help')) as g:

src/azdev/help.py

@@ -153,3 +153,7 @@
 helps['extension list'] = """
     short-summary: List what extensions are currently visible to your development environment.
 """
+
+helps['extension update-index'] = """
+    short-summary: Update the extensions index.json from a built WHL file.
+"""

src/azdev/operations/extensions.py → src/azdev/operations/extensions/__init__.py

@@ -1,8 +1,7 @@
-# -----------------------------------------------------------------------------
+# --------------------------------------------------------------------------------------------
 # Copyright (c) Microsoft Corporation. All rights reserved.
-# Licensed under the MIT License. See License.txt in the project root for
-# license information.
-# -----------------------------------------------------------------------------
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
 
 from collections import OrderedDict
 from glob import glob
@@ -85,3 +84,58 @@ def list_extensions():
         if long_name not in installed:
             results.append({'name': long_name, 'path': folder})
     return results
+
+
+def _get_sha256sum(a_file):
+    import hashlib
+    sha256 = hashlib.sha256()
+    with open(a_file, 'rb') as f:
+        sha256.update(f.read())
+    return sha256.hexdigest()
+
+
+def update_extension_index(extension):
+    import json
+    import re
+    import tempfile
+
+    from .util import get_ext_metadata, get_whl_from_url
+
+    NAME_REGEX = r'.*/([^/]*)-\d+.\d+.\d+'
+
+    ext_path = get_ext_repo_path()
+
+    # Get extension WHL from URL
+    if not extension.endswith('.whl') or not extension.startswith('https:'):
+        raise ValueError('usage error: only URL to a WHL file currently supported.')
+
+    # Extract the extension name
+    try:
+        extension_name = re.findall(NAME_REGEX, extension)[0]
+        extension_name = extension_name.replace('_', '-')
+    except IndexError:
+        raise ValueError('unable to parse extension name')
+
+    extensions_dir = tempfile.mkdtemp()
+    ext_dir = tempfile.mkdtemp(dir=extensions_dir)
+    whl_cache_dir = tempfile.mkdtemp()
+    whl_cache = {}
+    ext_file = get_whl_from_url(whl_path, extension_name, whl_cache_dir, whl_cache)
+
+    with open(os.join(ext_path, 'src', 'index.json'), 'r') as infile:
+        curr_index = json.loads(infile.read())
+
+    try:
+        entry = curr_index['extensions'][extension_name]
+    except IndexError:
+        raise ValueError('{} not found in index.json'.format(extension_name))
+
+    entry[0]['downloadUrl'] = whl_path
+    entry[0]['sha256Digest'] = _get_sha256sum(ext_file)
+    entry[0]['filename'] = whl_path.split('/')[-1]
+    entry[0]['metadata'] = get_ext_metadata(ext_dir, ext_file, extension_name)
+
+    # update index and write back to file
+    curr_index['extensions'][extension_name] = entry
+    with open(os.join(ext_path, 'src', 'index.json'), 'w') as outfile:
+        outfile.write(json.dumps(curr_index, indent=4, sort_keys=True))

src/azdev/operations/extensions/util.py

@@ -0,0 +1,71 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+import os
+import json
+import zipfile
+from wheel.install import WHEEL_INFO_RE
+
+from azdev.utilities import EXTENSION_PREFIX
+
+
+def _get_extension_modname(ext_dir):
+    # Modification of https://github.com/Azure/azure-cli/blob/dev/src/azure-cli-core/azure/cli/core/extension.py#L153
+    pos_mods = [n for n in os.listdir(ext_dir)
+                if n.startswith(EXTENSION_PREFIX) and os.path.isdir(os.path.join(ext_dir, n))]
+    if len(pos_mods) != 1:
+        raise AssertionError("Expected 1 module to load starting with "
+                             "'{}': got {}".format(EXTENSION_PREFIX, pos_mods))
+    return pos_mods[0]
+
+
+def _get_azext_metadata(ext_dir):
+    # Modification of https://github.com/Azure/azure-cli/blob/dev/src/azure-cli-core/azure/cli/core/extension.py#L109
+    AZEXT_METADATA_FILENAME = 'azext_metadata.json'
+    azext_metadata = None
+    ext_modname = _get_extension_modname(ext_dir=ext_dir)
+    azext_metadata_filepath = os.path.join(ext_dir, ext_modname, AZEXT_METADATA_FILENAME)
+    if os.path.isfile(azext_metadata_filepath):
+        with open(azext_metadata_filepath) as f:
+            azext_metadata = json.load(f)
+    return azext_metadata
+
+
+def get_ext_metadata(ext_dir, ext_file, ext_name):
+    # Modification of https://github.com/Azure/azure-cli/blob/dev/src/azure-cli-core/azure/cli/core/extension.py#L89
+    WHL_METADATA_FILENAME = 'metadata.json'
+    zip_ref = zipfile.ZipFile(ext_file, 'r')
+    zip_ref.extractall(ext_dir)
+    zip_ref.close()
+    metadata = {}
+    dist_info_dirs = [f for f in os.listdir(ext_dir) if f.endswith('.dist-info')]
+    azext_metadata = _get_azext_metadata(ext_dir)
+    if azext_metadata:
+        metadata.update(azext_metadata)
+    for dist_info_dirname in dist_info_dirs:
+        parsed_dist_info_dir = WHEEL_INFO_RE(dist_info_dirname)
+        if parsed_dist_info_dir and parsed_dist_info_dir.groupdict().get('name') == ext_name.replace('-', '_'):
+            whl_metadata_filepath = os.path.join(ext_dir, dist_info_dirname, WHL_METADATA_FILENAME)
+            if os.path.isfile(whl_metadata_filepath):
+                with open(whl_metadata_filepath) as f:
+                    metadata.update(json.load(f))
+    return metadata
+
+
+def get_whl_from_url(url, filename, tmp_dir, whl_cache=None):
+    if not whl_cache:
+        whl_cache = {}
+    if url in whl_cache:
+        return whl_cache[url]
+    import requests
+    r = requests.get(url, stream=True)
+    assert r.status_code == 200, "Request to {} failed with {}".format(url, r.status_code)
+    ext_file = os.path.join(tmp_dir, filename)
+    with open(ext_file, 'wb') as f:
+        for chunk in r.iter_content(chunk_size=1024):
+            if chunk:  # ignore keep-alive new chunks
+                f.write(chunk)
+    whl_cache[url] = ext_file
+    return ext_file

src/azdev/operations/performance.py

@@ -19,7 +19,6 @@
 logger = get_logger(__name__)
 
 TOTAL = 'ALL'
-NUM_RUNS = 3
 DEFAULT_THRESHOLD = 10
 
 # TODO: Treat everything as bubble instead of specific modules
@@ -33,14 +32,14 @@
 }
 
 
-def check_load_time():
+def check_load_time(runs=3):
     heading('Module Load Performance')
 
     regex = r"[^']*'([^']*)'[\D]*([\d\.]*)"
 
     results = {TOTAL: []}
     # Time the module loading X times
-    for i in range(0, NUM_RUNS + 1):
+    for i in range(0, runs + 1):
         lines = cmd('az -h --debug', show_stderr=True).result
         if i == 0:
             # Ignore the first run since it can be longer due to *.pyc file compilation

src/azdev/params.py

@@ -64,6 +64,12 @@ def load_arguments(self, command):
         c.argument('rules', options_list=['--rules', '-r'], nargs='+', help='Space-separated list of rules to run. Omit to run all rules.')
         c.argument('rule_types', options_list=['--rule-types', '-t'], nargs='+', choices=['params', 'commands', 'command_groups', 'help_entries'], help='Space-separated list of rule types to run. Omit to run all.')
 
+    with ArgumentsContext(self, 'perf') as c:
+        c.argument('runs', type=int, help='Number of runs to average performance over.')
+
     for scope in ['extension add', 'extension remove']:
         with ArgumentsContext(self, scope) as c:
             c.positional('extensions', metavar='NAME', nargs='+', help='Space-separated list of extension names.')
+
+    with ArgumentsContext(self, 'extension update-index') as c:
+        c.positional('extension', metavar='URL', help='URL to an extension WHL file.')

src/setup.py

@@ -45,6 +45,7 @@ def read(fname):
         'azdev.operations',
         'azdev.operations.linter',
         'azdev.operations.tests',
+        'azdev.operations.extensions',
         'azdev.utilities',
     ],
     install_requires=[