Add typing to ducktape

docs/conf.py

@@ -21,7 +21,7 @@
 import sphinx_rtd_theme
 from ducktape import __version__
 
-sys.path.insert(0, os.path.abspath('..'))
+sys.path.insert(0, os.path.abspath(".."))
 
 
 # -- General configuration ------------------------------------------------
@@ -33,24 +33,24 @@
 # Add any Sphinx extension module names here, as strings. They can be
 # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
 # ones.
-extensions = ['sphinx.ext.viewcode', 'sphinx.ext.autodoc', 'sphinxarg.ext']
+extensions = ["sphinx.ext.viewcode", "sphinx.ext.autodoc", "sphinxarg.ext"]
 
 # Add any paths that contain templates here, relative to this directory.
-templates_path = ['_templates']
+templates_path = ["_templates"]
 
 # The suffix(es) of source filenames.
 # You can specify multiple suffix as a list of string:
 #
 # source_suffix = ['.rst', '.md']
-source_suffix = '.rst'
+source_suffix = ".rst"
 
 # The master toctree document.
-master_doc = 'index'
+master_doc = "index"
 
 # General information about the project.
-project = u'Ducktape'
-copyright = u'2017, Confluent Inc.'
-author = u'Confluent Inc.'
+project = "Ducktape"
+copyright = "2017, Confluent Inc."
+author = "Confluent Inc."
 
 # The version info for the project you're documenting, acts as replacement for
 # |version| and |release|, also used in various other places throughout the
@@ -72,10 +72,10 @@
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
 # This patterns also effect to html_static_path and html_extra_path
-exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']
+exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
 
 # The name of the Pygments (syntax highlighting) style to use.
-pygments_style = 'sphinx'
+pygments_style = "sphinx"
 
 # If true, `todo` and `todoList` produce output, else they produce nothing.
 todo_include_todos = False
@@ -86,7 +86,7 @@
 # The theme to use for HTML and HTML Help pages.  See the documentation for
 # a list of builtin themes.
 #
-html_theme = 'sphinx_rtd_theme'
+html_theme = "sphinx_rtd_theme"
 
 # Theme options are theme-specific and customize the look and feel of a theme
 # further.  For a list of options available for each theme, see the
@@ -106,7 +106,7 @@
 # -- Options for HTMLHelp output ------------------------------------------
 
 # Output file base name for HTML help builder.
-htmlhelp_basename = 'ducktapedoc'
+htmlhelp_basename = "ducktapedoc"
 
 
 # -- Options for LaTeX output ---------------------------------------------
@@ -115,15 +115,12 @@
     # The paper size ('letterpaper' or 'a4paper').
     #
     # 'papersize': 'letterpaper',
-
     # The font size ('10pt', '11pt' or '12pt').
     #
     # 'pointsize': '10pt',
-
     # Additional stuff for the LaTeX preamble.
     #
     # 'preamble': '',
-
     # Latex figure (float) alignment
     #
     # 'figure_align': 'htbp',
@@ -133,19 +130,15 @@
 # (source start file, target name, title,
 #  author, documentclass [howto, manual, or own class]).
 latex_documents = [
-    (master_doc, 'ducktape.tex', u'Ducktape Documentation',
-     u'Confluent Inc.', 'manual'),
+    (master_doc, "ducktape.tex", "Ducktape Documentation", "Confluent Inc.", "manual"),
 ]
 
 
 # -- Options for manual page output ---------------------------------------
 
 # One entry per manual page. List of tuples
 # (source start file, name, description, authors, manual section).
-man_pages = [
-    (master_doc, 'ducktape', u'Ducktape Documentation',
-     [author], 1)
-]
+man_pages = [(master_doc, "ducktape", "Ducktape Documentation", [author], 1)]
 
 
 # -- Options for Texinfo output -------------------------------------------
@@ -154,15 +147,21 @@
 # (source start file, target name, title, author,
 #  dir menu entry, description, category)
 texinfo_documents = [
-    (master_doc, 'ducktape', u'Ducktape Documentation',
-     author, 'ducktape', 'One line description of project.',
-     'Miscellaneous'),
+    (
+        master_doc,
+        "ducktape",
+        "Ducktape Documentation",
+        author,
+        "ducktape",
+        "One line description of project.",
+        "Miscellaneous",
+    ),
 ]
 
 
 # ---- Options for Autodoc ------------------------------------------------
 
-autodoc_default_flags = ['show-inheritance']
+autodoc_default_flags = ["show-inheritance"]
 
 
 def skip(app, what, name, obj, skip, options):

ducktape/cluster/__init__.py

@@ -1,3 +1,3 @@
-from .localhost import LocalhostCluster  # NOQA
 from .json import JsonCluster  # NOQA
+from .localhost import LocalhostCluster  # NOQA
 from .vagrant import VagrantCluster  # NOQA

ducktape/cluster/cluster.py

@@ -13,37 +13,26 @@
 # limitations under the License.
 
 import collections
+from typing import Iterable, List, Union
 
-
-class ClusterNode(object):
-    def __init__(self, account, **kwargs):
-        self.account = account
-        for k, v in kwargs.items():
-            setattr(self, k, v)
-
-    @property
-    def name(self):
-        return self.account.hostname
-
-    @property
-    def operating_system(self):
-        return self.account.operating_system
+from ducktape.cluster.cluster_node import ClusterNode
+from ducktape.cluster.cluster_spec import ClusterSpec
 
 
 class Cluster(object):
-    """ Interface for a cluster -- a collection of nodes with login credentials.
+    """Interface for a cluster -- a collection of nodes with login credentials.
     This interface doesn't define any mapping of roles/services to nodes. It only interacts with some underlying
     system that can describe available resources and mediates reservations of those resources.
     """
 
     def __init__(self):
         self.max_used_nodes = 0
 
-    def __len__(self):
+    def __len__(self) -> int:
         """Size of this cluster object. I.e. number of 'nodes' in the cluster."""
         return self.available().size() + self.used().size()
 
-    def alloc(self, cluster_spec):
+    def alloc(self, cluster_spec) -> Union[ClusterNode, List[ClusterNode], "Cluster"]:
         """
         Allocate some nodes.
 
@@ -55,7 +44,7 @@ def alloc(self, cluster_spec):
         self.max_used_nodes = max(self.max_used_nodes, len(self.used()))
         return allocated
 
-    def do_alloc(self, cluster_spec):
+    def do_alloc(self, cluster_spec) -> Union[ClusterNode, List[ClusterNode], "Cluster"]:
         """
         Subclasses should implement actual allocation here.
 
@@ -65,15 +54,15 @@ def do_alloc(self, cluster_spec):
         """
         raise NotImplementedError
 
-    def free(self, nodes):
+    def free(self, nodes: Union[Iterable[ClusterNode], ClusterNode]) -> None:
         """Free the given node or list of nodes"""
         if isinstance(nodes, collections.abc.Iterable):
             for s in nodes:
                 self.free_single(s)
         else:
             self.free_single(nodes)
 
-    def free_single(self, node):
+    def free_single(self, node: ClusterNode) -> None:
         raise NotImplementedError()
 
     def __eq__(self, other):
@@ -82,22 +71,22 @@ def __eq__(self, other):
     def __hash__(self):
         return hash(tuple(sorted(self.__dict__.items())))
 
-    def num_available_nodes(self):
+    def num_available_nodes(self) -> int:
         return self.available().size()
 
-    def available(self):
+    def available(self) -> ClusterSpec:
         """
         Return a ClusterSpec object describing the currently available nodes.
         """
         raise NotImplementedError
 
-    def used(self):
+    def used(self) -> ClusterSpec:
         """
         Return a ClusterSpec object describing the currently in use nodes.
         """
         raise NotImplementedError
 
-    def max_used(self):
+    def max_used(self) -> int:
         return self.max_used_nodes
 
     def all(self):

ducktape/cluster/cluster_node.py

@@ -0,0 +1,18 @@
+from typing import Optional
+
+from ducktape.cluster.remoteaccount import RemoteAccount
+
+
+class ClusterNode(object):
+    def __init__(self, account: RemoteAccount, **kwargs):
+        self.account = account
+        for k, v in kwargs.items():
+            setattr(self, k, v)
+
+    @property
+    def name(self) -> Optional[str]:
+        return self.account.hostname
+
+    @property
+    def operating_system(self) -> Optional[str]:
+        return self.account.operating_system

ducktape/cluster/cluster_spec.py

@@ -15,39 +15,20 @@
 from __future__ import absolute_import
 
 import json
+import typing
 
 from ducktape.cluster.node_container import NodeContainer
 
-LINUX = "linux"
-
-WINDOWS = "windows"
-
-SUPPORTED_OS_TYPES = [LINUX, WINDOWS]
-
-
-class NodeSpec(object):
-    """
-    The specification for a ducktape cluster node.
-
-    :param operating_system:    The operating system of the node.
-    """
-    def __init__(self, operating_system=LINUX):
-        self.operating_system = operating_system
-        if self.operating_system not in SUPPORTED_OS_TYPES:
-            raise RuntimeError("Unsupported os type %s" % self.operating_system)
-
-    def __str__(self):
-        dict = {
-            "os": self.operating_system,
-        }
-        return json.dumps(dict, sort_keys=True)
+from .consts import LINUX
+from .node_spec import NodeSpec
 
 
 class ClusterSpec(object):
     """
     The specification for a ducktape cluster.
     """
-    nodes: NodeContainer = None
+
+    nodes: typing.Optional[NodeContainer] = None
 
     @staticmethod
     def empty():

ducktape/cluster/consts.py

@@ -0,0 +1,19 @@
+# Copyright 2017 Confluent Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+LINUX = "linux"
+
+WINDOWS = "windows"
+
+SUPPORTED_OS_TYPES = [LINUX, WINDOWS]

ducktape/cluster/finite_subcluster.py

@@ -11,23 +11,23 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import typing
 
-from ducktape.cluster.cluster import Cluster
+from ducktape.cluster.cluster import Cluster, ClusterNode
 from ducktape.cluster.cluster_spec import ClusterSpec
 from ducktape.cluster.node_container import NodeContainer
 
 
 class FiniteSubcluster(Cluster):
-    """This cluster class gives us a mechanism for allocating finite blocks of nodes from another cluster.
-    """
+    """This cluster class gives us a mechanism for allocating finite blocks of nodes from another cluster."""
 
-    def __init__(self, nodes):
+    def __init__(self, nodes: typing.Iterable[ClusterNode]):
         super(FiniteSubcluster, self).__init__()
         self.nodes = nodes
         self._available_nodes = NodeContainer(nodes)
         self._in_use_nodes = NodeContainer()
 
-    def do_alloc(self, cluster_spec):
+    def do_alloc(self, cluster_spec) -> typing.List[ClusterNode]:
         # there cannot be any bad nodes here,
         # since FiniteSubcluster operates on ClusterNode objects,
         # which are not checked for health by NodeContainer.remove_spec