enable and apply isort

Author: alexfikl

arraycontext/__init__.py

@@ -29,61 +29,38 @@
 """
 
 import sys
-from .context import (
-        ArrayContext,
-
-        Scalar, ScalarLike,
-        Array, ArrayT,
-        ArrayOrContainer, ArrayOrContainerT,
-        ArrayOrContainerOrScalar, ArrayOrContainerOrScalarT,
-
-        tag_axes)
-
-from .transform_metadata import (CommonSubexpressionTag,
-        ElementwiseMapKernelTag)
-
-# deprecated, remove in 2022.
-from .metadata import _FirstAxisIsElementsTag
 
 from .container import (
-        ArrayContainer, ArrayContainerT,
-        NotAnArrayContainerError,
-        is_array_container, is_array_container_type,
-        get_container_context_opt,
-        get_container_context_recursively, get_container_context_recursively_opt,
-        serialize_container, deserialize_container,
-        register_multivector_as_array_container)
+    ArrayContainer, ArrayContainerT, NotAnArrayContainerError, deserialize_container,
+    get_container_context_opt, get_container_context_recursively,
+    get_container_context_recursively_opt, is_array_container,
+    is_array_container_type, register_multivector_as_array_container,
+    serialize_container)
 from .container.arithmetic import with_container_arithmetic
 from .container.dataclass import dataclass_array_container
-
 from .container.traversal import (
-        map_array_container,
-        multimap_array_container,
-        rec_map_array_container,
-        rec_multimap_array_container,
-        mapped_over_array_containers,
-        multimapped_over_array_containers,
-        map_reduce_array_container,
-        multimap_reduce_array_container,
-        rec_map_reduce_array_container,
-        rec_multimap_reduce_array_container,
-        thaw, freeze,
-        flatten, unflatten, flat_size_and_dtype,
-        from_numpy, to_numpy,
-        outer, with_array_context)
-
-from .impl.pyopencl import PyOpenCLArrayContext
-from .impl.pytato import (PytatoPyOpenCLArrayContext,
-                          PytatoJAXArrayContext)
+    flat_size_and_dtype, flatten, freeze, from_numpy, map_array_container,
+    map_reduce_array_container, mapped_over_array_containers,
+    multimap_array_container, multimap_reduce_array_container,
+    multimapped_over_array_containers, outer, rec_map_array_container,
+    rec_map_reduce_array_container, rec_multimap_array_container,
+    rec_multimap_reduce_array_container, thaw, to_numpy, unflatten,
+    with_array_context)
+from .context import (
+    Array, ArrayContext, ArrayOrContainer, ArrayOrContainerOrScalar,
+    ArrayOrContainerOrScalarT, ArrayOrContainerT, ArrayT, Scalar, ScalarLike,
+    tag_axes)
 from .impl.jax import EagerJAXArrayContext
-
-from .pytest import (
-        PytestArrayContextFactory,
-        PytestPyOpenCLArrayContextFactory,
-        pytest_generate_tests_for_array_contexts,
-        pytest_generate_tests_for_pyopencl_array_context)
-
+from .impl.pyopencl import PyOpenCLArrayContext
+from .impl.pytato import PytatoJAXArrayContext, PytatoPyOpenCLArrayContext
 from .loopy import make_loopy_program
+# deprecated, remove in 2022.
+from .metadata import _FirstAxisIsElementsTag
+from .pytest import (
+    PytestArrayContextFactory, PytestPyOpenCLArrayContextFactory,
+    pytest_generate_tests_for_array_contexts,
+    pytest_generate_tests_for_pyopencl_array_context)
+from .transform_metadata import CommonSubexpressionTag, ElementwiseMapKernelTag
 
 
 __all__ = (

arraycontext/container/__init__.py

@@ -69,17 +69,19 @@
 """
 
 from functools import singledispatch
-from arraycontext.context import ArrayContext
-from typing import Any, Iterable, Tuple, Optional, TypeVar, Protocol, TYPE_CHECKING
-import numpy as np
+from typing import TYPE_CHECKING, Any, Iterable, Optional, Protocol, Tuple, TypeVar
 
 # For use in singledispatch type annotations, because sphinx can't figure out
 # what 'np' is.
 import numpy
+import numpy as np
+
+from arraycontext.context import ArrayContext
 
 
 if TYPE_CHECKING:
     from pymbolic.geometric_algebra import MultiVector
+
     from arraycontext import ArrayOrContainer
 
 

arraycontext/container/arithmetic.py

@@ -7,6 +7,7 @@
 
 import enum
 
+
 __copyright__ = """
 Copyright (C) 2020-1 University of Illinois Board of Trustees
 """
@@ -31,8 +32,8 @@
 THE SOFTWARE.
 """
 
+from typing import Any, Callable, Optional, Tuple, Type, TypeVar, Union
 from warnings import warn
-from typing import Any, Callable, Optional, Tuple, TypeVar, Union, Type
 
 import numpy as np
 

arraycontext/container/dataclass.py

@@ -30,9 +30,9 @@
 THE SOFTWARE.
 """
 
+from dataclasses import Field, fields, is_dataclass
 from typing import Tuple, Union, get_args, get_origin
 
-from dataclasses import Field, is_dataclass, fields
 from arraycontext.container import is_array_container_type
 
 
@@ -95,7 +95,7 @@ def is_array_field(f: Field) -> bool:
             # NOTE:
             # * `_BaseGenericAlias` catches `List`, `Tuple`, etc.
             # * `_SpecialForm` catches `Any`, `Literal`, etc.
-            from typing import (                    # type: ignore[attr-defined]
+            from typing import (  # type: ignore[attr-defined]
                 _BaseGenericAlias, _SpecialForm)
             if isinstance(f.type, (_BaseGenericAlias, _SpecialForm)):
                 # NOTE: anything except a Union is not allowed

arraycontext/container/traversal.py

@@ -41,6 +41,7 @@
 
 from __future__ import annotations
 
+
 __copyright__ = """
 Copyright (C) 2020-1 University of Illinois Board of Trustees
 """
@@ -65,22 +66,18 @@
 THE SOFTWARE.
 """
 
-from typing import Any, Callable, Iterable, List, Optional, Union, Tuple, cast
-from functools import update_wrapper, partial, singledispatch
+from functools import partial, singledispatch, update_wrapper
+from typing import Any, Callable, Iterable, List, Optional, Tuple, Union, cast
 from warnings import warn
 
 import numpy as np
 
-from arraycontext.context import (
-    ArrayT, ArrayOrContainer, ArrayOrContainerT,
-    ArrayOrContainerOrScalar, ScalarLike,
-    ArrayContext, Array
-)
 from arraycontext.container import (
-        NotAnArrayContainerError,
-        ArrayContainer,
-        serialize_container, deserialize_container,
-        get_container_context_recursively_opt)
+    ArrayContainer, NotAnArrayContainerError, deserialize_container,
+    get_container_context_recursively_opt, serialize_container)
+from arraycontext.context import (
+    Array, ArrayContext, ArrayOrContainer, ArrayOrContainerOrScalar,
+    ArrayOrContainerT, ArrayT, ScalarLike)
 
 
 # {{{ array container traversal helpers

arraycontext/context.py

@@ -160,15 +160,18 @@
 
 from abc import ABC, abstractmethod
 from typing import (
-        Any, Callable, Dict, Optional, Tuple, Union, Mapping, Protocol, TypeVar,
-        TYPE_CHECKING)
+    TYPE_CHECKING, Any, Callable, Dict, Mapping, Optional, Protocol, Tuple, TypeVar,
+    Union)
 
 import numpy as np
+
 from pytools import memoize_method
 from pytools.tag import ToTagSetConvertible
 
+
 if TYPE_CHECKING:
     import loopy
+
     from arraycontext.container import ArrayContainer
 
 
@@ -426,8 +429,9 @@ def _get_einsum_prg(self,
                         spec: str, arg_names: Tuple[str, ...],
                         tagged: ToTagSetConvertible) -> "loopy.TranslationUnit":
         import loopy as lp
-        from .loopy import _DEFAULT_LOOPY_OPTIONS
         from loopy.version import MOST_RECENT_LANGUAGE_VERSION
+
+        from .loopy import _DEFAULT_LOOPY_OPTIONS
         return lp.make_einsum(
             spec,
             arg_names,

arraycontext/fake_numpy.py

@@ -24,6 +24,7 @@
 
 
 import numpy as np
+
 from arraycontext.container import NotAnArrayContainerError, serialize_container
 from arraycontext.container.traversal import rec_map_array_container
 
@@ -105,8 +106,8 @@ def conjugate(self, x):
 # {{{ BaseFakeNumpyLinalgNamespace
 
 def _reduce_norm(actx, arys, ord):
-    from numbers import Number
     from functools import reduce
+    from numbers import Number
 
     if ord is None:
         ord = 2

arraycontext/impl/jax/__init__.py

@@ -32,9 +32,10 @@
 import numpy as np
 
 from pytools.tag import ToTagSetConvertible
-from arraycontext.context import ArrayContext, Array, ArrayOrContainer, ScalarLike
-from arraycontext.container.traversal import (with_array_context,
-                                              rec_map_array_container)
+
+from arraycontext.container.traversal import (
+    rec_map_array_container, with_array_context)
+from arraycontext.context import Array, ArrayContext, ArrayOrContainer, ScalarLike
 
 
 class EagerJAXArrayContext(ArrayContext):

arraycontext/impl/jax/fake_numpy.py

@@ -23,17 +23,15 @@
 """
 from functools import partial, reduce
 
-import numpy as np
 import jax.numpy as jnp
+import numpy as np
 
-from arraycontext.fake_numpy import (
-        BaseFakeNumpyNamespace, BaseFakeNumpyLinalgNamespace,
-        )
-from arraycontext.container.traversal import (
-        rec_multimap_array_container, rec_map_array_container,
-        rec_map_reduce_array_container,
-        )
 from arraycontext.container import NotAnArrayContainerError, serialize_container
+from arraycontext.container.traversal import (
+    rec_map_array_container, rec_map_reduce_array_container,
+    rec_multimap_array_container)
+from arraycontext.fake_numpy import (
+    BaseFakeNumpyLinalgNamespace, BaseFakeNumpyNamespace)
 
 
 class EagerJAXFakeNumpyLinalgNamespace(BaseFakeNumpyLinalgNamespace):

arraycontext/impl/pyopencl/__init__.py

@@ -28,21 +28,21 @@
 THE SOFTWARE.
 """
 
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple
 from warnings import warn
-from typing import Callable, Dict, List, Optional, Tuple, TYPE_CHECKING
 
 import numpy as np
 
 from pytools.tag import ToTagSetConvertible
 
-from arraycontext.context import ArrayContext, Array, ArrayOrContainer, ScalarLike
-from arraycontext.container.traversal import (rec_map_array_container,
-                                              with_array_context)
+from arraycontext.container.traversal import (
+    rec_map_array_container, with_array_context)
+from arraycontext.context import Array, ArrayContext, ArrayOrContainer, ScalarLike
 
 
 if TYPE_CHECKING:
-    import pyopencl
     import loopy as lp
+    import pyopencl
 
 
 # {{{ PyOpenCLArrayContext
@@ -287,6 +287,7 @@ def call_loopy(self, t_unit, **kwargs):
                 wait_event_queue.pop(0).wait()
 
         import arraycontext.impl.pyopencl.taggable_cl_array as tga
+
         # FIXME: Inherit loopy tags for these arrays
         return {name: tga.to_tagged_cl_array(ary) for name, ary in result.items()}