diff --git a/requirements_travis.txt b/requirements_travis.txt index 5bc90a4bc..ac073ab34 100644 --- a/requirements_travis.txt +++ b/requirements_travis.txt @@ -1,5 +1,5 @@ jax jaxlib -pytype +pytype>=2019.06.21 pytest -torch>=1.1 \ No newline at end of file +torch>=1.1 diff --git a/tensornetwork/network_components.py b/tensornetwork/network_components.py index da3e1f3c0..e0ddfb7bf 100644 --- a/tensornetwork/network_components.py +++ b/tensornetwork/network_components.py @@ -16,7 +16,8 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function -from typing import Any, Dict, List, Optional, Set, Text, Tuple, Type, Union +from typing import Any, Dict, List, Optional, Set, Text, Tuple, Type, Union, \ + overload import typing import numpy as np from tensornetwork.backends import base_backend @@ -256,7 +257,18 @@ def has_nondangling_edge(self): return True return False + @overload + def __getitem__(self, key: slice) -> List["Edge"]: + pass + + @overload def __getitem__(self, key: Union[int, Text]) -> "Edge": + pass + + def __getitem__(self, key: Union[int, Text, slice] + ) -> Union["Edge", List["Edge"]]: + if isinstance(key, slice): + return self.edges[key] return self.get_edge(key) def __str__(self) -> Text: diff --git a/tensornetwork/network_components_test.py b/tensornetwork/network_components_test.py index 88a51ff1c..ddbb1b02a 100644 --- a/tensornetwork/network_components_test.py +++ b/tensornetwork/network_components_test.py @@ -201,8 +201,17 @@ def test_node_magic_getitem(single_node_edge): edge = single_node_edge.edge node.add_edge(edge, axis=0) assert node[0] == edge - - + + +def test_node_magic_getslice(single_node_edge): + node = single_node_edge.node + edge = single_node_edge.edge + node.add_edge(edge, axis=0) + assert node[:1][0] is edge + assert len(node[None:None]) == 3 + assert len(node[0:3:2]) == 2 + + def test_node_magic_str(single_node_edge): node = single_node_edge.node assert str(node) == node.name