[Typing][C-116,C-117] Add type annotations for "paddle/geometric/*" (#66792)

successfulbarrier · web-flow · commit d365726a0bc5 · 2024-08-01T18:00:58.000+08:00
* test

* 添加chunk图例

* 删除测试

* pre-commit完善

* 1

* 将文档修改部分还原

* 修改并添加了neighbors的类型注释

* overload重载

* 修正overload

* 再次修改

* 再次修改2
diff --git a/python/paddle/geometric/reindex.py b/python/paddle/geometric/reindex.py
@@ -12,19 +12,31 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Sequence
+
 import paddle
 from paddle import _C_ops
 from paddle.base.data_feeder import check_variable_and_dtype
 from paddle.base.framework import Variable
 from paddle.base.layer_helper import LayerHelper
 from paddle.framework import in_dynamic_or_pir_mode
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+
 __all__ = []
 
 
 def reindex_graph(
-    x, neighbors, count, value_buffer=None, index_buffer=None, name=None
-):
+    x: Tensor,
+    neighbors: Tensor,
+    count: Tensor,
+    value_buffer: Tensor | None = None,
+    index_buffer: Tensor | None = None,
+    name: str | None = None,
+) -> tuple[Tensor, Tensor, Tensor]:
     """
 
     Reindex Graph API.
@@ -137,8 +149,13 @@ def reindex_graph(
 
 
 def reindex_heter_graph(
-    x, neighbors, count, value_buffer=None, index_buffer=None, name=None
-):
+    x: Tensor,
+    neighbors: Sequence[Tensor],
+    count: Sequence[Tensor],
+    value_buffer: Tensor | None = None,
+    index_buffer: Tensor | None = None,
+    name: str | None = None,
+) -> tuple[Tensor, Tensor, Tensor]:
     """
 
     Reindex HeterGraph API.
diff --git a/python/paddle/geometric/sampling/neighbors.py b/python/paddle/geometric/sampling/neighbors.py
@@ -12,14 +12,62 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal, overload
+
 from paddle import _C_ops
 from paddle.base.data_feeder import check_variable_and_dtype
 from paddle.base.layer_helper import LayerHelper
 from paddle.framework import in_dynamic_or_pir_mode
 
+if TYPE_CHECKING:
+    from paddle import Tensor
 __all__ = []
 
 
+@overload
+def sample_neighbors(
+    row: Tensor,
+    colptr: Tensor,
+    input_nodes: Tensor,
+    sample_size: int = ...,
+    eids: Tensor | None = ...,
+    return_eids: Literal[True] = ...,
+    perm_buffer: Tensor | None = ...,
+    name: str | None = ...,
+) -> tuple[Tensor, Tensor, Tensor]:
+    ...
+
+
+@overload
+def sample_neighbors(
+    row: Tensor,
+    colptr: Tensor,
+    input_nodes: Tensor,
+    sample_size: int = ...,
+    eids: Tensor | None = ...,
+    return_eids: Literal[False] = ...,
+    perm_buffer: Tensor | None = ...,
+    name: str | None = ...,
+) -> tuple[Tensor, Tensor]:
+    ...
+
+
+@overload
+def sample_neighbors(
+    row: Tensor,
+    colptr: Tensor,
+    input_nodes: Tensor,
+    sample_size: int = ...,
+    eids: Tensor | None = ...,
+    return_eids: bool = ...,
+    perm_buffer: Tensor | None = ...,
+    name: str | None = ...,
+) -> tuple[Tensor, Tensor] | tuple[Tensor, Tensor, Tensor]:
+    ...
+
+
 def sample_neighbors(
     row,
     colptr,
@@ -169,6 +217,48 @@ def sample_neighbors(
     return out_neighbors, out_count
 
 
+@overload
+def weighted_sample_neighbors(
+    row: Tensor,
+    colptr: Tensor,
+    edge_weight: Tensor,
+    input_nodes: Tensor,
+    sample_size: int = ...,
+    eids: Tensor | None = ...,
+    return_eids: Literal[True] = ...,
+    name: str | None = ...,
+) -> tuple[Tensor, Tensor, Tensor]:
+    ...
+
+
+@overload
+def weighted_sample_neighbors(
+    row: Tensor,
+    colptr: Tensor,
+    edge_weight: Tensor,
+    input_nodes: Tensor,
+    sample_size: int = ...,
+    eids: Tensor | None = ...,
+    return_eids: Literal[False] = ...,
+    name: str | None = ...,
+) -> tuple[Tensor, Tensor]:
+    ...
+
+
+@overload
+def weighted_sample_neighbors(
+    row: Tensor,
+    colptr: Tensor,
+    edge_weight: Tensor,
+    input_nodes: Tensor,
+    sample_size: int = ...,
+    eids: Tensor | None = ...,
+    return_eids: bool = ...,
+    name: str | None = ...,
+) -> tuple[Tensor, Tensor] | tuple[Tensor, Tensor, Tensor]:
+    ...
+
+
 def weighted_sample_neighbors(
     row,
     colptr,
