Bug 1507860: [taskgraph] Move most clases to use attrs; r=dustin

This moves most of the low-hanging fruit to use attrs.

Differential Revision: https://phabricator.services.mozilla.com/D1141

--HG--
extra : moz-landing-system : lando

taskcluster/taskgraph/generator.py

@@ -7,6 +7,7 @@ import logging
 import os
 import yaml
 import copy
+import attr
 
 from . import filter_tasks
 from .graph import Graph
@@ -20,7 +21,7 @@ from .util.verify import (
     verify_docs,
     verifications,
 )
-from .config import load_graph_config
+from .config import load_graph_config, GraphConfig
 
 logger = logging.getLogger(__name__)
 
@@ -31,13 +32,13 @@ class KindNotFound(Exception):
     """
 
 
+@attr.s(frozen=True)
 class Kind(object):
 
-    def __init__(self, name, path, config, graph_config):
-        self.name = name
-        self.path = path
-        self.config = config
-        self.graph_config = graph_config
+    name = attr.ib(type=basestring)
+    path = attr.ib(type=basestring)
+    config = attr.ib(type=dict)
+    graph_config = attr.ib(type=GraphConfig)
 
     def _get_loader(self):
         try:

taskcluster/taskgraph/graph.py

@@ -4,9 +4,11 @@
 
 from __future__ import absolute_import, print_function, unicode_literals
 
+import attr
 import collections
 
 
+@attr.s(frozen=True)
 class Graph(object):
     """
     Generic representation of a directed acyclic graph with labeled edges
@@ -23,22 +25,8 @@ class Graph(object):
     node `left` to node `right..
     """
 
-    def __init__(self, nodes, edges):
-        """
-        Create a graph.  Nodes and edges are both as described in the class
-        documentation.  Both values are used by reference, and should not be
-        modified after building a graph.
-        """
-        assert isinstance(nodes, set)
-        assert isinstance(edges, set)
-        self.nodes = nodes
-        self.edges = edges
-
-    def __eq__(self, other):
-        return self.nodes == other.nodes and self.edges == other.edges
-
-    def __repr__(self):
-        return "<Graph nodes={!r} edges={!r}>".format(self.nodes, self.edges)
+    nodes = attr.ib(converter=frozenset)
+    edges = attr.ib(converter=frozenset)
 
     def transitive_closure(self, nodes, reverse=False):
         """

taskcluster/taskgraph/task.py

@@ -4,7 +4,10 @@
 
 from __future__ import absolute_import, print_function, unicode_literals
 
+import attr
 
+
+@attr.s
 class Task(object):
     """
     Representation of a task in a TaskGraph.  Each Task has, at creation:
@@ -24,40 +27,21 @@ class Task(object):
     This class is just a convenience wrapper for the data type and managing
     display, comparison, serialization, etc. It has no functionality of its own.
     """
-    def __init__(self, kind, label, attributes, task,
-                 optimization=None, dependencies=None,
-                 release_artifacts=None):
-        self.kind = kind
-        self.label = label
-        self.attributes = attributes
-        self.task = task
 
-        self.task_id = None
+    kind = attr.ib()
+    label = attr.ib()
+    attributes = attr.ib()
+    task = attr.ib()
+    task_id = attr.ib(default=None, init=False)
+    optimization = attr.ib(default=None)
+    dependencies = attr.ib(factory=dict)
+    release_artifacts = attr.ib(
+        converter=attr.converters.optional(frozenset),
+        default=None,
+    )
 
-        self.attributes['kind'] = kind
-
-        self.optimization = optimization
-        self.dependencies = dependencies or {}
-        if release_artifacts:
-            self.release_artifacts = frozenset(release_artifacts)
-        else:
-            self.release_artifacts = None
-
-    def __eq__(self, other):
-        return self.kind == other.kind and \
-            self.label == other.label and \
-            self.attributes == other.attributes and \
-            self.task == other.task and \
-            self.task_id == other.task_id and \
-            self.optimization == other.optimization and \
-            self.dependencies == other.dependencies and \
-            self.release_artifacts == other.release_artifacts
-
-    def __repr__(self):
-        return ('Task({kind!r}, {label!r}, {attributes!r}, {task!r}, '
-                'optimization={optimization!r}, '
-                'dependencies={dependencies!r}, '
-                'release_artifacts={release_artifacts!r})'.format(**self.__dict__))
+    def __attrs_post_init__(self):
+        self.attributes['kind'] = self.kind
 
     def to_json(self):
         rv = {
@@ -71,7 +55,7 @@ class Task(object):
         if self.task_id:
             rv['task_id'] = self.task_id
         if self.release_artifacts:
-            rv['release_artifacts'] = sorted(self.release_artifacts),
+            rv['release_artifacts'] = sorted(self.release_artifacts)
         return rv
 
     @classmethod
@@ -88,7 +72,7 @@ class Task(object):
             task=task_dict['task'],
             optimization=task_dict['optimization'],
             dependencies=task_dict.get('dependencies'),
-            release_artifacts=task_dict.get('release-artifacts')
+            release_artifacts=task_dict.get('release-artifacts'),
         )
         if 'task_id' in task_dict:
             rv.task_id = task_dict['task_id']

taskcluster/taskgraph/taskgraph.py

@@ -7,7 +7,10 @@ from __future__ import absolute_import, print_function, unicode_literals
 from .graph import Graph
 from .task import Task
 
+import attr
 
+
+@attr.s(frozen=True)
 class TaskGraph(object):
     """
     Representation of a task graph.
@@ -16,10 +19,11 @@ class TaskGraph(object):
     by label. TaskGraph instances should be treated as immutable.
     """
 
-    def __init__(self, tasks, graph):
-        assert set(tasks) == graph.nodes
-        self.tasks = tasks
-        self.graph = graph
+    tasks = attr.ib()
+    graph = attr.ib()
+
+    def __attrs_post_init__(self):
+        assert set(self.tasks) == self.graph.nodes
 
     def for_each_task(self, f, *args, **kwargs):
         for task_label in self.graph.visit_postorder():
@@ -37,12 +41,6 @@ class TaskGraph(object):
         "Iterate over tasks in undefined order"
         return self.tasks.itervalues()
 
-    def __repr__(self):
-        return "<TaskGraph graph={!r} tasks={!r}>".format(self.graph, self.tasks)
-
-    def __eq__(self, other):
-        return self.tasks == other.tasks and self.graph == other.graph
-
     def to_json(self):
         "Return a JSON-able object representing the task graph, as documented"
         named_links_dict = self.graph.named_links_dict()

taskcluster/taskgraph/transforms/base.py

@@ -4,33 +4,40 @@
 
 from __future__ import absolute_import, print_function, unicode_literals
 
+import attr
 
+from ..parameters import Parameters
+from ..config import GraphConfig
+
+
+@attr.s(frozen=True)
 class TransformConfig(object):
-    """A container for configuration affecting transforms.  The `config`
-    argument to transforms is an instance of this class, possibly with
-    additional kind-specific attributes beyond those set here."""
-    def __init__(self, kind, path, config, params,
-                 kind_dependencies_tasks=None, graph_config=None):
-        # the name of the current kind
-        self.kind = kind
+    """
+    A container for configuration affecting transforms.  The `config` argument
+    to transforms is an instance of this class.
+    """
 
-        # the path to the kind configuration directory
-        self.path = path
+    # the name of the current kind
+    kind = attr.ib()
 
-        # the parsed contents of kind.yml
-        self.config = config
+    # the path to the kind configuration directory
+    path = attr.ib(type=basestring)
 
-        # the parameters for this task-graph generation run
-        self.params = params
+    # the parsed contents of kind.yml
+    config = attr.ib(type=dict)
 
-        # a list of all the tasks associated with the kind dependencies of the
-        # current kind
-        self.kind_dependencies_tasks = kind_dependencies_tasks
+    # the parameters for this task-graph generation run
+    params = attr.ib(type=Parameters)
 
-        # Global configuration of the taskgraph
-        self.graph_config = graph_config or {}
+    # a list of all the tasks associated with the kind dependencies of the
+    # current kind
+    kind_dependencies_tasks = attr.ib()
+
+    # Global configuration of the taskgraph
+    graph_config = attr.ib(type=GraphConfig)
 
 
+@attr.s()
 class TransformSequence(object):
     """
     Container for a sequence of transforms.  Each transform is represented as a
@@ -42,22 +49,15 @@ class TransformSequence(object):
     sequence.
     """
 
-    def __init__(self, transforms=None):
-        self.transforms = transforms or []
+    _transforms = attr.ib(factory=list)
 
     def __call__(self, config, items):
-        for xform in self.transforms:
+        for xform in self._transforms:
             items = xform(config, items)
             if items is None:
                 raise Exception("Transform {} is not a generator".format(xform))
         return items
 
-    def __repr__(self):
-        return '\n'.join(
-            ['TransformSequence(['] +
-            [repr(x) for x in self.transforms] +
-            ['])'])
-
     def add(self, func):
-        self.transforms.append(func)
+        self._transforms.append(func)
         return func

taskcluster/taskgraph/util/seta.py

@@ -10,6 +10,7 @@ import requests
 from collections import defaultdict
 from redo import retry
 from requests import exceptions
+import attr
 
 logger = logging.getLogger(__name__)
 
@@ -23,19 +24,20 @@ SETA_ENDPOINT = "https://treeherder.mozilla.org/api/project/%s/seta/" \
 PUSH_ENDPOINT = "https://hg.mozilla.org/integration/%s/json-pushes/?startID=%d&endID=%d"
 
 
+@attr.s(frozen=True)
 class SETA(object):
     """
     Interface to the SETA service, which defines low-value tasks that can be optimized out
     of the taskgraph.
     """
-    def __init__(self):
-        # cached low value tasks, by project
-        self.low_value_tasks = {}
-        self.low_value_bb_tasks = {}
-        # cached push dates by project
-        self.push_dates = defaultdict(dict)
-        # cached push_ids that failed to retrieve datetime for
-        self.failed_json_push_calls = []
+
+    # cached low value tasks, by project
+    low_value_tasks = attr.ib(factory=dict, init=False)
+    low_value_bb_tasks = attr.ib(factory=dict, init=False)
+    # cached push dates by project
+    push_dates = attr.ib(factory=lambda: defaultdict(dict), init=False)
+    # cached push_ids that failed to retrieve datetime for
+    failed_json_push_calls = attr.ib(factory=list, init=False)
 
     def _get_task_string(self, task_tuple):
         # convert task tuple to single task string, so the task label sent in can match

taskcluster/taskgraph/util/verify.py

@@ -10,12 +10,15 @@ import re
 import os
 import sys
 
+import attr
+
 from .. import GECKO
 
 logger = logging.getLogger(__name__)
 base_path = os.path.join(GECKO, 'taskcluster', 'docs')
 
 
+@attr.s(frozen=True)
 class VerificationSequence(object):
     """
     Container for a sequence of verifications over a TaskGraph. Each
@@ -24,11 +27,10 @@ class VerificationSequence(object):
     time with no task but with the taskgraph and the same scratch_pad
     that was passed for each task.
     """
-    def __init__(self):
-        self.verifications = {}
+    _verifications = attr.ib(factory=dict)
 
     def __call__(self, graph_name, graph):
-        for verification in self.verifications.get(graph_name, []):
+        for verification in self._verifications.get(graph_name, []):
             scratch_pad = {}
             graph.for_each_task(verification, scratch_pad=scratch_pad)
             verification(None, graph, scratch_pad=scratch_pad)
@@ -36,7 +38,7 @@ class VerificationSequence(object):
 
     def add(self, graph_name):
         def wrap(func):
-            self.verifications.setdefault(graph_name, []).append(func)
+            self._verifications.setdefault(graph_name, []).append(func)
             return func
         return wrap