Major improvements to spack create

lib/spack/spack/cmd/checksum.py

@@ -22,6 +22,8 @@
 # License along with this program; if not, write to the Free Software
 # Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
 ##############################################################################
+from __future__ import print_function
+
 import argparse
 import hashlib
 
@@ -30,93 +32,133 @@
 import spack.cmd
 import spack.util.crypto
 from spack.stage import Stage, FailedDownloadError
+from spack.util.naming import *
 from spack.version import *
 
 description = "Checksum available versions of a package."
 
 
 def setup_parser(subparser):
     subparser.add_argument(
-        'package', metavar='PACKAGE', help='Package to list versions for')
+        'package',
+        help='Package to checksum versions for')
     subparser.add_argument(
-        '--keep-stage', action='store_true', dest='keep_stage',
+        '--keep-stage', action='store_true',
         help="Don't clean up staging area when command completes.")
     subparser.add_argument(
         'versions', nargs=argparse.REMAINDER,
         help='Versions to generate checksums for')
 
 
-def get_checksums(versions, urls, **kwargs):
-    # Allow commands like create() to do some analysis on the first
-    # archive after it is downloaded.
+def get_checksums(url_dict, name, **kwargs):
+    """Fetches and checksums archives from URLs.
+
+    This function is called by both ``spack checksum`` and ``spack create``.
+    The ``first_stage_function`` kwarg allows ``spack create`` to determine
+    things like the build system of the archive.
+
+    :param dict url_dict: A dictionary of the form: version -> URL
+    :param str name: The name of the package
+    :param callable first_stage_function: Function to run on first staging area
+    :param bool keep_stage: Don't clean up staging area when command completes
+
+    :returns: A multi-line string containing versions and corresponding hashes
+    :rtype: str
+    """
     first_stage_function = kwargs.get('first_stage_function', None)
     keep_stage = kwargs.get('keep_stage', False)
 
+    sorted_versions = sorted(url_dict.keys(), reverse=True)
+
+    # Find length of longest string in the list for padding
+    max_len = max(len(str(v)) for v in sorted_versions)
+    num_ver = len(sorted_versions)
+
+    tty.msg("Found {0} version{1} of {2}:".format(
+            num_ver, '' if num_ver == 1 else 's', name),
+            "",
+            *spack.cmd.elide_list(
+                ["{0:{1}}  {2}".format(v, max_len, url_dict[v])
+                 for v in sorted_versions]))
+    print()
+
+    archives_to_fetch = tty.get_number(
+        "How many would you like to checksum?", default=1, abort='q')
+
+    if not archives_to_fetch:
+        tty.die("Aborted.")
+
+    versions = sorted_versions[:archives_to_fetch]
+    urls = [url_dict[v] for v in versions]
+
     tty.msg("Downloading...")
-    hashes = []
+    version_hashes = []
     i = 0
     for url, version in zip(urls, versions):
         try:
             with Stage(url, keep=keep_stage) as stage:
+                # Fetch the archive
                 stage.fetch()
                 if i == 0 and first_stage_function:
+                    # Only run first_stage_function the first time,
+                    # no need to run it every time
                     first_stage_function(stage, url)
 
-                hashes.append((version, spack.util.crypto.checksum(
+                # Checksum the archive and add it to the list
+                version_hashes.append((version, spack.util.crypto.checksum(
                     hashlib.md5, stage.archive_file)))
                 i += 1
-        except FailedDownloadError as e:
-            tty.msg("Failed to fetch %s" % url)
+        except FailedDownloadError:
+            tty.msg("Failed to fetch {0}".format(url))
         except Exception as e:
-            tty.msg('Something failed on %s, skipping.\n    (%s)' % (url, e))
+            tty.msg("Something failed on {0}, skipping.".format(url),
+                    "  ({0})".format(e))
 
-    return hashes
+    if not version_hashes:
+        tty.die("Could not fetch any versions for {0}".format(name))
+
+    # Find length of longest string in the list for padding
+    max_len = max(len(str(v)) for v, h in version_hashes)
+
+    # Generate the version directives to put in a package.py
+    version_lines = "\n".join([
+        "    version('{0}', {1}'{2}')".format(
+            v, ' ' * (max_len - len(str(v))), h) for v, h in version_hashes
+    ])
+
+    num_hash = len(version_hashes)
+    tty.msg("Checksummed {0} version{1} of {2}".format(
+        num_hash, '' if num_hash == 1 else 's', name))
+
+    return version_lines
 
 
 def checksum(parser, args):
-    # get the package we're going to generate checksums for
+    # Make sure the user provided a package and not a URL
+    if not valid_fully_qualified_module_name(args.package):
+        tty.die("`spack checksum` accepts package names, not URLs. "
+                "Use `spack md5 <url>` instead.")
+
+    # Get the package we're going to generate checksums for
     pkg = spack.repo.get(args.package)
 
-    # If the user asked for specific versions, use those.
     if args.versions:
-        versions = {}
+        # If the user asked for specific versions, use those
+        url_dict = {}
         for version in args.versions:
             version = ver(version)
             if not isinstance(version, Version):
-                tty.die("Cannot generate checksums for version lists or " +
-                        "version ranges.  Use unambiguous versions.")
-            versions[version] = pkg.url_for_version(version)
+                tty.die("Cannot generate checksums for version lists or "
+                        "version ranges. Use unambiguous versions.")
+            url_dict[version] = pkg.url_for_version(version)
     else:
-        versions = pkg.fetch_remote_versions()
-        if not versions:
-            tty.die("Could not fetch any versions for %s" % pkg.name)
-
-    sorted_versions = sorted(versions, reverse=True)
-
-    # Find length of longest string in the list for padding
-    maxlen = max(len(str(v)) for v in versions)
+        # Otherwise, see what versions we can find online
+        url_dict = pkg.fetch_remote_versions()
+        if not url_dict:
+            tty.die("Could not find any versions for {0}".format(pkg.name))
 
-    tty.msg("Found %s versions of %s" % (len(versions), pkg.name),
-            *spack.cmd.elide_list(
-                ["{0:{1}}  {2}".format(v, maxlen, versions[v])
-                 for v in sorted_versions]))
-    print
-    archives_to_fetch = tty.get_number(
-        "How many would you like to checksum?", default=5, abort='q')
-
-    if not archives_to_fetch:
-        tty.msg("Aborted.")
-        return
-
-    version_hashes = get_checksums(
-        sorted_versions[:archives_to_fetch],
-        [versions[v] for v in sorted_versions[:archives_to_fetch]],
-        keep_stage=args.keep_stage)
-
-    if not version_hashes:
-        tty.die("Could not fetch any versions for %s" % pkg.name)
+    version_lines = get_checksums(
+        url_dict, pkg.name, keep_stage=args.keep_stage)
 
-    version_lines = [
-        "  version('%s', '%s')" % (v, h) for v, h in version_hashes
-    ]
-    tty.msg("Checksummed new versions of %s:" % pkg.name, *version_lines)
+    print()
+    print(version_lines)