[Refactor] Do not expose internal function in field, exception, expr, any_array , _ndrange, _ndarray

python/taichi/lang/__init__.py

@@ -16,12 +16,9 @@
 from taichi._lib import core as _ti_core
 from taichi._lib.utils import locale_encode
 from taichi.lang import impl
-from taichi.lang._ndarray import ScalarNdarray
-from taichi.lang._ndrange import GroupedNDRange, ndrange
-from taichi.lang.any_array import AnyArray, AnyArrayAccess
+from taichi.lang._ndrange import ndrange
 from taichi.lang.enums import Layout
-from taichi.lang.exception import (InvalidOperationError,
-                                   TaichiCompilationError, TaichiNameError,
+from taichi.lang.exception import (TaichiCompilationError, TaichiNameError,
                                    TaichiSyntaxError, TaichiTypeError)
 from taichi.lang.expr import Expr, make_expr_group
 from taichi.lang.field import Field, ScalarField
@@ -716,17 +713,19 @@ def cache_read_only(*args):
 
 def assume_in_range(val, base, low, high):
     return _ti_core.expr_assume_in_range(
-        Expr(val).ptr,
-        Expr(base).ptr, low, high)
+        impl.Expr(val).ptr,
+        impl.Expr(base).ptr, low, high)
 
 
 def loop_unique(val, covers=None):
     if covers is None:
         covers = []
     if not isinstance(covers, (list, tuple)):
         covers = [covers]
-    covers = [x.snode.ptr if isinstance(x, Expr) else x.ptr for x in covers]
-    return _ti_core.expr_loop_unique(Expr(val).ptr, covers)
+    covers = [
+        x.snode.ptr if isinstance(x, impl.Expr) else x.ptr for x in covers
+    ]
+    return _ti_core.expr_loop_unique(impl.Expr(val).ptr, covers)
 
 
 parallelize = _ti_core.parallelize

python/taichi/lang/_ndarray.py

@@ -354,3 +354,6 @@ def setter(value):
 
         self.getter = getter
         self.setter = setter
+
+
+__all__ = []

python/taichi/lang/_ndrange.py

@@ -4,7 +4,7 @@
 from taichi.lang.matrix import _IntermediateMatrix
 
 
-class ndrange:
+class _Ndrange:
     def __init__(self, *args):
         args = list(args)
         for i, arg in enumerate(args):
@@ -42,10 +42,17 @@ def grouped(self):
         return GroupedNDRange(self)
 
 
+def ndrange(*args):
+    return _Ndrange(*args)
+
+
 class GroupedNDRange:
     def __init__(self, r):
         self.r = r
 
     def __iter__(self):
         for ind in self.r:
             yield _IntermediateMatrix(len(ind), 1, list(ind))
+
+
+__all__ = ['ndrange']

python/taichi/lang/any_array.py

@@ -69,3 +69,6 @@ def subscript(self, i, j):
             indices = self.indices_first + indices_second
         return Expr(_ti_core.subscript(self.arr.ptr,
                                        make_expr_group(*indices)))
+
+
+__all__ = []

python/taichi/lang/ast/ast_transformer.py

@@ -8,7 +8,7 @@
 from taichi._lib import core as _ti_core
 from taichi.lang import expr, impl, kernel_arguments, kernel_impl, matrix, mesh
 from taichi.lang import ops as ti_ops
-from taichi.lang._ndrange import ndrange
+from taichi.lang._ndrange import _Ndrange, ndrange
 from taichi.lang.ast.ast_transformer_utils import Builder, LoopStatus
 from taichi.lang.ast.symbol_resolver import ASTResolver
 from taichi.lang.exception import TaichiSyntaxError
@@ -731,7 +731,7 @@ def build_static_for(ctx, node, is_grouped):
         if is_grouped:
             assert len(node.iter.args[0].args) == 1
             ndrange_arg = build_stmt(ctx, node.iter.args[0].args[0])
-            if not isinstance(ndrange_arg, ndrange):
+            if not isinstance(ndrange_arg, _Ndrange):
                 raise TaichiSyntaxError(
                     "Only 'ti.ndrange' is allowed in 'ti.static(ti.grouped(...))'."
                 )

python/taichi/lang/enums.py

@@ -9,3 +9,6 @@ class Layout(Enum):
     """
     AOS = 1
     SOA = 2
+
+
+__all__ = ['Layout']

python/taichi/lang/exception.py

@@ -25,3 +25,9 @@ def handle_exception_from_cpp(exc):
     if isinstance(exc, core.TaichiTypeError):
         return TaichiTypeError(str(exc))
     return exc
+
+
+__all__ = [
+    'TaichiSyntaxError', 'TaichiTypeError', 'TaichiCompilationError',
+    'TaichiNameError'
+]

python/taichi/lang/expr.py

@@ -75,3 +75,6 @@ def make_expr_group(*exprs):
         else:
             expr_group.push_back(Expr(i).ptr)
     return expr_group
+
+
+__all__ = []

python/taichi/lang/field.py

@@ -314,3 +314,6 @@ class SNodeHostAccess:
     def __init__(self, accessor, key):
         self.accessor = accessor
         self.key = key
+
+
+__all__ = []

python/taichi/lang/impl.py

@@ -6,7 +6,7 @@
 from taichi._lib import core as _ti_core
 from taichi._logging import warn
 from taichi.lang._ndarray import ScalarNdarray
-from taichi.lang._ndrange import GroupedNDRange, ndrange
+from taichi.lang._ndrange import GroupedNDRange, _Ndrange
 from taichi.lang.any_array import AnyArray, AnyArrayAccess
 from taichi.lang.exception import InvalidOperationError, TaichiTypeError
 from taichi.lang.expr import Expr, make_expr_group
@@ -51,7 +51,7 @@ def expr_init(rhs):
         return rhs
     if isinstance(rhs, _ti_core.Arch):
         return rhs
-    if isinstance(rhs, ndrange):
+    if isinstance(rhs, _Ndrange):
         return rhs
     if isinstance(rhs, MeshElementFieldProxy):
         return rhs
@@ -848,8 +848,9 @@ def static(x, *xs):
     if len(xs):  # for python-ish pointer assign: x, y = ti.static(y, x)
         return [static(x)] + [static(x) for x in xs]
 
-    if isinstance(x, (bool, int, float, range, list, tuple, enumerate, ndrange,
-                      GroupedNDRange, zip, filter, map)) or x is None:
+    if isinstance(x,
+                  (bool, int, float, range, list, tuple, enumerate, _Ndrange,
+                   GroupedNDRange, zip, filter, map)) or x is None:
         return x
     if isinstance(x, AnyArray):
         return x
@@ -874,7 +875,7 @@ def grouped(x):
         >>> for I in ti.grouped(ndrange(8, 16)):
         >>>     print(I[0] + I[1])
     """
-    if isinstance(x, ndrange):
+    if isinstance(x, _Ndrange):
         return x.grouped()
     return x
 

python/taichi/linalg/sparse_solver.py

@@ -1,6 +1,7 @@
 import numpy as np
 import taichi.lang
 from taichi._lib import core as _ti_core
+from taichi.lang.field import Field
 from taichi.linalg import SparseMatrix
 from taichi.types.primitive_types import f32
 
@@ -70,7 +71,7 @@ def solve(self, b):
         Returns:
             numpy.array: The solution of linear systems.
         """
-        if isinstance(b, taichi.lang.Field):
+        if isinstance(b, Field):
             return self.solver.solve(b.to_numpy())
         if isinstance(b, np.ndarray):
             return self.solver.solve(b)