From 04bebab004b5c53f76cc4aa04c26e4adf36a3f9e Mon Sep 17 00:00:00 2001
From: HansBug <killog@126.com>
Date: Mon, 20 Sep 2021 20:39:16 +0800
Subject: [PATCH] doc(hansbug): add documentation for treetensor.common and
 treetensor.utils

---
 docs/source/api_doc/common/trees.rst |  9 ++++-
 docs/source/api_doc/numpy/index.rst  |  4 +++
 docs/source/api_doc/torch/index.rst  |  4 +++
 docs/source/api_doc/utils/doc.rst    | 14 ++++++++
 treetensor/common/object.py          |  5 +++
 treetensor/common/trees.py           | 53 +++++++++++++++++++++++++++-
 6 files changed, 87 insertions(+), 2 deletions(-)
 create mode 100644 docs/source/api_doc/utils/doc.rst

diff --git a/docs/source/api_doc/common/trees.rst b/docs/source/api_doc/common/trees.rst
index 43778fa15..8888a5cda 100644
--- a/docs/source/api_doc/common/trees.rst
+++ b/docs/source/api_doc/common/trees.rst
@@ -13,5 +13,12 @@ BaseTreeStruct
 -----------------------
 
 .. autoclass:: BaseTreeStruct
-    :members:
+    :members: __repr__, __str__
+
+
+clsmeta
+------------------------
+
+.. autofunction:: clsmeta
+
 
diff --git a/docs/source/api_doc/numpy/index.rst b/docs/source/api_doc/numpy/index.rst
index ec4936623..c8616ec39 100644
--- a/docs/source/api_doc/numpy/index.rst
+++ b/docs/source/api_doc/numpy/index.rst
@@ -1,6 +1,10 @@
 treetensor.numpy
 =====================
 
+.. py:currentmodule:: treetensor.numpy
+
+.. automodule:: treetensor.numpy
+
 .. toctree::
     :maxdepth: 2
 
diff --git a/docs/source/api_doc/torch/index.rst b/docs/source/api_doc/torch/index.rst
index d26aa1152..6ea87158c 100644
--- a/docs/source/api_doc/torch/index.rst
+++ b/docs/source/api_doc/torch/index.rst
@@ -1,6 +1,10 @@
 treetensor.torch
 =====================
 
+.. py:currentmodule:: treetensor.torch
+
+.. automodule:: treetensor.torch
+
 .. toctree::
     :maxdepth: 2
 
diff --git a/docs/source/api_doc/utils/doc.rst b/docs/source/api_doc/utils/doc.rst
new file mode 100644
index 000000000..dc9b97934
--- /dev/null
+++ b/docs/source/api_doc/utils/doc.rst
@@ -0,0 +1,14 @@
+Documentation Decorators
+=============================
+
+.. py:currentmodule:: treetensor.utils
+
+.. automodule:: treetensor.utils.doc
+
+
+doc_from
+------------------
+
+.. autofunction:: doc_from
+
+
diff --git a/treetensor/common/object.py b/treetensor/common/object.py
index b7c0a170f..f16271f79 100644
--- a/treetensor/common/object.py
+++ b/treetensor/common/object.py
@@ -10,6 +10,11 @@ def _object(obj):
 
 
 class Object(BaseTreeStruct, metaclass=clsmeta(_object, allow_dict=True)):
+    """
+    Overview:
+        Generic object tree class, used in :py:mod:`treetensor.numpy` and :py:mod:`treetensor.torch`.
+    """
+
     def __init__(self, data):
         """
         In :class:`treetensor.common.Object`, object or object tree can be initialized.
diff --git a/treetensor/common/trees.py b/treetensor/common/trees.py
index 8e3e874a4..48c98cbd8 100644
--- a/treetensor/common/trees.py
+++ b/treetensor/common/trees.py
@@ -3,6 +3,7 @@ import io
 import os
 from functools import partial
 from typing import Optional, Tuple, Callable
+from typing import Type
 
 from treevalue import func_treelize as original_func_treelize
 from treevalue import general_tree_value, TreeValue
@@ -20,6 +21,17 @@ __all__ = [
 
 def print_tree(tree: TreeValue, repr_: Callable = str,
                ascii_: bool = False, show_node_id: bool = True, file=None):
+    """
+    Overview:
+        Print a tree structure to the given file.
+
+    Arguments:
+        - tree (:obj:`TreeValue`): Given tree object.
+        - repr\_ (:obj:`Callable`): Representation function, default is ``str``.
+        - ascii\_ (:obj:`bool`): Use ascii to print the tree, default is ``False``.
+        - show_node_id (:obj:`bool`): Show node id of the tree, default is ``True``.
+        - file: Output file of this print procedure, default is ``None`` which means to stdout. 
+    """
     print_to_file = partial(builtins.print, file=file)
     node_ids = {}
     if ascii_:
@@ -86,15 +98,54 @@ class BaseTreeStruct(general_tree_value()):
     """
 
     def __repr__(self):
+        """
+        Get the tree-based representation format of this object.
+
+        Examples::
+
+            >>> from treetensor.common import Object
+            >>> repr(Object(1))  # Object is subclass of BaseTreeStruct
+            '1'
+
+            >>> repr(Object({'a': 1, 'b': 2, 'x': {'c': 3, 'd': 4}}))
+            '<Object 0x7fe00b121220>\n├── a --> 1\n├── b --> 2\n└── x --> <Object 0x7fe00b121c10>\n    ├── c --> 3\n    └── d --> 4\n'
+
+            >>> Object({'a': 1, 'b': 2, 'x': {'c': 3, 'd': 4}})
+            <Object 0x7fe00b1271c0>
+            ├── a --> 1
+            ├── b --> 2
+            └── x --> <Object 0x7fe00b127910>
+                ├── c --> 3
+                └── d --> 4
+        """
         with io.StringIO() as sfile:
             print_tree(self, repr_=repr, ascii_=False, file=sfile)
             return sfile.getvalue()
 
     def __str__(self):
+        """
+        The same as :py:meth:`BaseTreeStruct.__repr__`.
+        """
         return self.__repr__()
 
 
-def clsmeta(func, allow_dict: bool = False):
+def clsmeta(func, allow_dict: bool = False) -> Type[type]:
+    """
+    Overview:
+        Create a metaclass based on generating function.
+
+        Used in :py:class:`treetensor.common.Object`,
+        :py:class:`treetensor.torch.Tensor` and :py:class:`treetensor.torch.Size`.
+        Can do modify onto the constructor function of the classes.
+
+    Arguments:
+        - func: Generating function.
+        - allow_dict (:obj:`bool`): Auto transform dictionary to :py:class:`treevalue.TreeValue` class, \
+                                    default is ``False``.
+    Returns:
+        - metaclass (:obj:`Type[type]`): Metaclass for creating a new class.
+    """
+
     class _TempTreeValue(TreeValue):
         pass
 
-- 
GitLab