Bug 1387135 - Add ability to use 'artifact' template to |mach try fuzzy|, r=maja_zf

This uses the new templating mechanism in taskgraph to schedule artifact builds when using
|mach try fuzzy|. Like |mach try syntax|, this will automatically be enabled if
--enable-artifact-builds is set in the mozconfig. The --artifact/--no-artifact arguments
can be used to override the default.

MozReview-Commit-ID: J8HVZzOt4mX

--HG--
extra : rebase_source : 130fa85ed24ef8f0fe5a713f254bbffd1a83b5a7

tools/tryselect/mach_commands.py

@@ -43,6 +43,11 @@ def syntax_parser():
     return parser
 
 
+def fuzzy_parser():
+    from tryselect.selectors.fuzzy import FuzzyParser
+    return FuzzyParser()
+
+
 @CommandProvider
 class TrySelect(MachCommandBase):
 
@@ -69,19 +74,8 @@ class TrySelect(MachCommandBase):
 
     @SubCommand('try',
                 'fuzzy',
-                description='Select tasks on try using a fuzzy finder')
-    @CommandArgument('-q', '--query', metavar='STR',
-                     help="Use the given query instead of entering the selection "
-                          "interface. Equivalent to typing <query><ctrl-a><enter> "
-                          "from the interface.")
-    @CommandArgument('-u', '--update', action='store_true', default=False,
-                     help="Update fzf before running.")
-    @CommandArgument('--full', action='store_true', default=False,
-                     help="Use the full set of tasks as input to fzf (instead of "
-                          "target tasks).")
-    @CommandArgument('-p', '--parameters', default=None,
-                     help="Use the given parameters.yml to generate tasks, "
-                          "defaults to latest parameters.yml from mozilla-central")
+                description='Select tasks on try using a fuzzy finder',
+                parser=fuzzy_parser)
     def try_fuzzy(self, *args, **kwargs):
         """Select which tasks to use with fzf.
 

tools/tryselect/selectors/fuzzy.py

@@ -13,6 +13,7 @@ from distutils.spawn import find_executable
 from mozboot.util import get_state_dir
 
 from ..tasks import generate_tasks
+from ..templates import TemplateParser
 from ..vcs import VCSHelper
 
 try:
@@ -80,6 +81,35 @@ fzf_header_shortcuts = {
 }
 
 
+class FuzzyParser(TemplateParser):
+    arguments = [
+        [['-q', '--query'],
+         {'metavar': 'STR',
+          'help': "Use the given query instead of entering the selection "
+                  "interface. Equivalent to typing <query><ctrl-a><enter> "
+                  "from the interface.",
+          }],
+        [['-u', '--update'],
+         {'action': 'store_true',
+          'default': False,
+          'help': "Update fzf before running.",
+          }],
+        [['--full'],
+         {'action': 'store_true',
+          'default': False,
+          'help': "Use the full set of tasks as input to fzf (instead of "
+                  "target tasks).",
+          }],
+        [['-p', '--parameters'],
+         {'default': None,
+          'help': "Use the given parameters.yml to generate tasks, "
+                  "defaults to latest parameters.yml from mozilla-central",
+          }],
+    ]
+
+    templates = ['artifact']
+
+
 def run(cmd, cwd=None):
     is_win = platform.system() == 'Windows'
     return subprocess.call(cmd, cwd=cwd, shell=True if is_win else False)
@@ -166,7 +196,7 @@ def format_header():
     return FZF_HEADER.format(shortcuts=', '.join(shortcuts), t=terminal)
 
 
-def run_fuzzy_try(update=False, query=None, full=False, parameters=None):
+def run_fuzzy_try(update=False, query=None, templates=None, full=False, parameters=None, **kwargs):
     fzf = fzf_bootstrap(update)
 
     if not fzf:
@@ -199,4 +229,5 @@ def run_fuzzy_try(update=False, query=None, full=False, parameters=None):
         print("no tasks selected")
         return
 
-    return vcs.push_to_try("Pushed via 'mach try fuzzy', see diff for scheduled tasks", selected)
+    msg = "Pushed via 'mach try fuzzy', see diff for scheduled tasks"
+    return vcs.push_to_try(msg, selected, templates)

tools/tryselect/templates.py

@@ -0,0 +1,96 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+"""
+Templates provide a way of modifying the task definition of selected
+tasks. They live under taskcluster/taskgraph/templates.
+"""
+
+from __future__ import absolute_import, print_function, unicode_literals
+
+import os
+from abc import ABCMeta, abstractmethod, abstractproperty
+from argparse import ArgumentParser
+
+from mozbuild.base import BuildEnvironmentNotFoundException, MozbuildObject
+
+here = os.path.abspath(os.path.dirname(__file__))
+
+
+class TemplateParser(ArgumentParser):
+    __metaclass__ = ABCMeta
+
+    @abstractproperty
+    def arguments(self):
+        pass
+
+    @abstractproperty
+    def templates(self):
+        pass
+
+    def __init__(self, *args, **kwargs):
+        ArgumentParser.__init__(self, *args, **kwargs)
+
+        for cli, kwargs in self.arguments:
+            self.add_argument(*cli, **kwargs)
+
+        self.templates = {t: all_templates[t]() for t in self.templates}
+
+        group = self.add_argument_group("template arguments")
+        for template in self.templates.values():
+            template.add_arguments(group)
+
+    def parse_known_args(self, *args, **kwargs):
+        args, remainder = ArgumentParser.parse_known_args(self, *args, **kwargs)
+
+        if self.templates:
+            args.templates = {}
+            for name, cls in self.templates.iteritems():
+                context = cls.context(**vars(args))
+                if context is not None:
+                    args.templates[name] = context
+
+        return args, remainder
+
+
+class Template(object):
+    __metaclass__ = ABCMeta
+
+    @abstractmethod
+    def add_arguments(self, parser):
+        pass
+
+    @abstractmethod
+    def context(self, **kwargs):
+        pass
+
+
+class Artifact(Template):
+
+    def add_arguments(self, parser):
+        group = parser.add_mutually_exclusive_group()
+        group.add_argument('--artifact', action='store_true',
+                           help='Force artifact builds where possible.')
+        group.add_argument('--no-artifact', action='store_true',
+                           help='Disable artifact builds even if being used locally.')
+
+    def context(self, artifact, no_artifact, **kwargs):
+        if artifact:
+            return {'enabled': 1}
+
+        if no_artifact:
+            return
+
+        build = MozbuildObject.from_environment(cwd=here)
+        try:
+            if build.substs.get("MOZ_ARTIFACT_BUILDS"):
+                print("Artifact builds enabled, pass --no-artifact to disable")
+                return {'enabled': 1}
+        except BuildEnvironmentNotFoundException:
+            pass
+
+
+all_templates = {
+    'artifact': Artifact,
+}

tools/tryselect/vcs.py

@@ -76,10 +76,14 @@ class VCSHelper(object):
             print(e.output)
             raise
 
-    def write_task_config(self, labels):
+    def write_task_config(self, labels, templates=None):
         config = os.path.join(self.root, 'try_task_config.json')
         with open(config, 'w') as fh:
-            json.dump(sorted(labels), fh, indent=2)
+            try_task_config = {'tasks': sorted(labels)}
+            if templates:
+                try_task_config['templates'] = templates
+
+            json.dump(try_task_config, fh, indent=2)
         return config
 
     def check_working_directory(self):
@@ -102,11 +106,11 @@ class VCSHelper(object):
 
 class HgHelper(VCSHelper):
 
-    def push_to_try(self, msg, labels=None):
+    def push_to_try(self, msg, labels=None, templates=None):
         self.check_working_directory()
 
         if labels:
-            config = self.write_task_config(labels)
+            config = self.write_task_config(labels, templates)
             self.run(['hg', 'add', config])
 
         try:
@@ -136,7 +140,7 @@ class HgHelper(VCSHelper):
 
 class GitHelper(VCSHelper):
 
-    def push_to_try(self, msg, labels=None):
+    def push_to_try(self, msg, labels=None, templates=None):
         self.check_working_directory()
 
         try:
@@ -146,7 +150,7 @@ class GitHelper(VCSHelper):
             return 1
 
         if labels:
-            config = self.write_task_config(labels)
+            config = self.write_task_config(labels, templates)
             self.run(['git', 'add', config])
 
         subprocess.check_call(['git', 'commit', '--allow-empty', '-m', msg])