Removed six

Sonnet 2 does not support 2.X, so it makes little sense to keep six around.

PiperOrigin-RevId: 380235488

requirements.txt

@@ -1,6 +1,5 @@
 absl-py>=0.7.1
 numpy>=1.16.3
-six>=1.12.0
 dm-tree>=0.1.1
 wrapt>=1.11.1
 tabulate>=0.7.5

sonnet/src/BUILD

@@ -11,7 +11,6 @@ snt_py_library(
         ":once",
         ":types",
         ":utils",
-        # pip: six
         # pip: tensorflow
     ],
 )
@@ -24,7 +23,6 @@ snt_py_test(
         ":test_utils",
         # pip: absl/testing:parameterized
         # pip: numpy
-        # pip: six
         # pip: tensorflow
         # pip: wrapt
     ],
@@ -239,7 +237,6 @@ snt_py_library(
     deps = [
         ":types",
         # pip: numpy
-        # pip: six
         # pip: tensorflow
     ],
 )
@@ -333,7 +330,6 @@ snt_py_library(
     deps = [
         ":base",
         ":once",
-        # pip: six
         # pip: tensorflow
     ],
 )
@@ -375,7 +371,6 @@ snt_py_library(
         ":once",
         ":types",
         ":utils",
-        # pip: six
         # pip: tensorflow
         # pip: tree
     ],
@@ -400,7 +395,6 @@ snt_py_library(
     srcs = ["regularizers.py"],
     deps = [
         ":types",
-        # pip: six
         # pip: tensorflow
     ],
 )
@@ -513,7 +507,6 @@ snt_py_library(
     deps = [
         ":initializers",
         # pip: absl/logging
-        # pip: six
         # pip: tabulate
         # pip: tensorflow
     ],

sonnet/src/base.py

@@ -22,7 +22,6 @@
 import sys
 from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Type, TypeVar
 
-import six
 from sonnet.src import once
 from sonnet.src import types
 from sonnet.src import utils
@@ -181,7 +180,7 @@ def auto_repr(cls: Type[Any], *args, **kwargs) -> str:
   Returns:
     A string representing a call equivalent to `cls(*args, **kwargs)`.
   """
-  argspec = utils.getfullargspec(cls.__init__)
+  argspec = inspect.getfullargspec(cls.__init__)
   arg_names = argspec.args
   # Keep used positionals minus self.
   arg_names = arg_names[1:(len(args) + 1)]
@@ -255,7 +254,7 @@ def wrap_with_name_scope(
         "The super constructor must be called before any other methods in "
         "your constructor. If this is not possible then annotate all the "
         "methods called with `@snt.no_name_scope`.")
-    six.raise_from(exc_value, exc_value_from)
+    raise exc_value from exc_value_from
 
   with module_name_scope:
     # snt.Module enters the module name scope for all methods. To disable this
@@ -366,7 +365,7 @@ def assert_tf2():
 assert_tf2.checked = False
 
 
-class Module(six.with_metaclass(ModuleMetaclass, tf.Module)):
+class Module(tf.Module, metaclass=ModuleMetaclass):
   """Base class for Sonnet modules.
 
   A Sonnet module is a lightweight container for variables and other modules.

sonnet/src/base_test.py

@@ -18,7 +18,6 @@
 
 from absl.testing import parameterized
 import numpy as np
-import six
 from sonnet.src import base
 from sonnet.src import test_utils
 import tensorflow as tf
@@ -409,8 +408,7 @@ def __init__(self, depth, trainable=True):
     self.w = tf.Variable(1.0, trainable=trainable, name="mushroom")
 
 
-@six.add_metaclass(abc.ABCMeta)
-class AbstractModule(base.Module):
+class AbstractModule(base.Module, metaclass=abc.ABCMeta):
 
   @abc.abstractmethod
   def __call__(self, x):
@@ -632,8 +630,7 @@ def __new__(cls, name, bases, clsdict):
     return new_type
 
 
-@six.add_metaclass(CombiningMetaclass)
-class ModuleWithCustomMetaclass(base.Module):
+class ModuleWithCustomMetaclass(base.Module, metaclass=CombiningMetaclass):
 
   def __init__(self):
     super(ModuleWithCustomMetaclass, self).__init__()

sonnet/src/conformance/BUILD

@@ -13,7 +13,6 @@ snt_py_library(
     deps = [
         # pip: absl/testing:parameterized
         # pip: numpy
-        # pip: six
         "//sonnet",
         # pip: tensorflow
     ],
@@ -23,7 +22,6 @@ snt_py_test(
     name = "api_test",
     srcs = ["api_test.py"],
     deps = [
-        # pip: six
         "//sonnet",
         "//sonnet/src:test_utils",
         # pip: tensorflow

sonnet/src/conformance/api_test.py

@@ -14,7 +14,8 @@
 # ============================================================================
 """Tests for Sonnet's public API."""
 
-from six.moves import reload_module as reload
+import importlib
+
 import sonnet as snt
 from sonnet.src import test_utils
 import tensorflow as tf
@@ -28,7 +29,7 @@ def test_src_not_exported(self):
   def test_supports_reload(self):
     mysnt = snt
     for _ in range(2):
-      mysnt = reload(mysnt)
+      mysnt = importlib.reload(mysnt)
       self.assertFalse(hasattr(mysnt, "src"))
 
 

sonnet/src/conformance/goldens.py

@@ -19,7 +19,6 @@
 
 from absl.testing import parameterized
 import numpy as np
-import six
 import sonnet as snt
 import tensorflow as tf
 
@@ -75,8 +74,7 @@ def range_like(t, start=0):
       t.shape)
 
 
-@six.add_metaclass(abc.ABCMeta)
-class Golden:
+class Golden(abc.ABC):
   """Represents a golden checkpoint file."""
 
   @abc.abstractmethod
@@ -95,7 +93,6 @@ def forward(self, module, x=None):
     pass
 
 
-@six.add_metaclass(abc.ABCMeta)
 class AbstractGolden(Golden):
   """Abstract base class for golden tests of single input modules."""
 

sonnet/src/initializers.py

@@ -18,13 +18,11 @@
 import collections
 from typing import Iterable, Mapping, Optional, Union
 import numpy as np
-import six
 from sonnet.src import types
 import tensorflow as tf
 
 
-@six.add_metaclass(abc.ABCMeta)
-class Initializer:
+class Initializer(abc.ABC):
   """Initializer base class, all initializers must implement a call method."""
 
   @abc.abstractmethod

sonnet/src/metrics.py

@@ -17,14 +17,12 @@
 import abc
 from typing import Optional
 
-import six
 from sonnet.src import base
 from sonnet.src import once
 import tensorflow as tf
 
 
-@six.add_metaclass(abc.ABCMeta)
-class Metric(base.Module):
+class Metric(base.Module, metaclass=abc.ABCMeta):
   """Metric base class."""
 
   @abc.abstractmethod

sonnet/src/optimizers/BUILD

@@ -73,7 +73,6 @@ snt_py_library(
     srcs = ["rmsprop.py"],
     deps = [
         ":optimizer_utils",
-        # pip: six
         "//sonnet/src:base",
         "//sonnet/src:once",
         "//sonnet/src:types",

sonnet/src/optimizers/rmsprop.py

@@ -14,9 +14,9 @@
 # ============================================================================
 """RMSProp module."""
 
+import itertools
 from typing import Optional, Sequence, Union
 
-import six
 from sonnet.src import base
 from sonnet.src import once
 from sonnet.src import types
@@ -127,7 +127,7 @@ def apply(self, updates: Sequence[types.ParameterUpdate],
     optimizer_utils.check_distribution_strategy()
     optimizer_utils.check_updates_parameters(updates, parameters)
     self._initialize(parameters)
-    for update, parameter, mom_var, ms_var, mg_var in six.moves.zip_longest(
+    for update, parameter, mom_var, ms_var, mg_var in itertools.zip_longest(
         updates, parameters, self.mom, self.ms, self.mg):
       if update is None:
         continue

sonnet/src/recurrent.py

@@ -20,7 +20,6 @@
 from typing import Optional, Sequence, Tuple, Union
 import uuid
 
-import six
 from sonnet.src import base
 from sonnet.src import conv
 from sonnet.src import initializers
@@ -39,8 +38,7 @@
 # pylint: enable=g-direct-tensorflow-import
 
 
-@six.add_metaclass(abc.ABCMeta)
-class RNNCore(base.Module):
+class RNNCore(base.Module, metaclass=abc.ABCMeta):
   """Base class for Recurrent Neural Network cores.
 
   This class defines the basic functionality that every core should
@@ -84,8 +82,7 @@ def initial_state(self, batch_size: types.IntegerLike, **kwargs):
     """
 
 
-@six.add_metaclass(abc.ABCMeta)
-class UnrolledRNN(base.Module):
+class UnrolledRNN(base.Module, metaclass=abc.ABCMeta):
   """Base class for unrolled Recurrent Neural Networks.
 
   This class is a generalization of :class:`RNNCore` which operates on
@@ -261,7 +258,7 @@ def static_unroll(
   outputs = None
   state = initial_state
   output_accs = None
-  for t in six.moves.range(num_steps):
+  for t in range(num_steps):
     outputs, state = _rnn_step(
         core,
         input_tas,

sonnet/src/regularizers.py

@@ -17,13 +17,11 @@
 import abc
 from typing import Sequence
 
-import six
 from sonnet.src import types
 import tensorflow as tf
 
 
-@six.add_metaclass(abc.ABCMeta)
-class Regularizer:
+class Regularizer(abc.ABC):
   """Base regularizer class."""
 
   @abc.abstractmethod

sonnet/src/utils.py

@@ -21,7 +21,6 @@
 from typing import Any, Callable, Dict, Generic, Optional, Sequence, Tuple, TypeVar, Union
 
 from absl import logging
-import six
 from sonnet.src import initializers
 import tabulate
 import tensorflow as tf
@@ -74,7 +73,7 @@ def _decorate_bound_method(*args, **kwargs):
 
       return _decorate_bound_method
 
-    argspec = getfullargspec(f)
+    argspec = inspect.getfullargspec(f)
     if argspec.args and argspec.args[0] == "self":
 
       @functools.wraps(f)
@@ -183,16 +182,6 @@ def smart_autograph_wrapper(*args, **kwargs):
   return smart_autograph_wrapper
 
 
-def getfullargspec(func):
-  """Gets the names and default values of a function's parameters."""
-  if six.PY2:
-    # Assume that we are running with PyType patched Python 2.7 and getargspec
-    # will not barf if `func` has type annotations.
-    return inspect.getargspec(func)  # pylint: disable=deprecated-method
-  else:
-    return inspect.getfullargspec(func)
-
-
 def variable_like(inputs: Union[tf.Tensor, tf.Variable],
                   initializer: initializers.Initializer = initializers.Zeros(),
                   trainable: Optional[bool] = None,