From 2839e884aa7c9dae85eba3589a2bfe0e16cafea7 Mon Sep 17 00:00:00 2001
From: Hao Zhang <hzhangxyz@outlook.com>
Date: Wed, 22 Nov 2023 23:24:14 +0800
Subject: [PATCH] Update Edge type output format.

---
 tat/edge.py        | 9 +++++++--
 tests/test_edge.py | 4 ++--
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/tat/edge.py b/tat/edge.py
index 68d2c8e7a..be34aa174 100644
--- a/tat/edge.py
+++ b/tat/edge.py
@@ -197,10 +197,15 @@ def __str__(self: Edge) -> str:
         # pylint: disable=no-else-return
         if any(self.fermion):
             # Fermionic edge
-            return f"(dimension={self.dimension}, arrow={self.arrow}, fermion={self.fermion}, symmetry={self.symmetry})"
+            fermion = ','.join(str(sub_fermion) for sub_fermion in self.fermion)
+            symmetry = ','.join(
+                f"[{','.join(str(sub_sym.item()) for sub_sym in sub_symmetry)}]" for sub_symmetry in self.symmetry)
+            return f"(dimension={self.dimension}, arrow={self.arrow}, fermion=({fermion}), symmetry=({symmetry}))"
         elif self.fermion:
             # Bosonic edge
-            return f"(dimension={self.dimension}, symmetry={self.symmetry})"
+            symmetry = ','.join(
+                f"[{','.join(str(sub_sym.item()) for sub_sym in sub_symmetry)}]" for sub_symmetry in self.symmetry)
+            return f"(dimension={self.dimension}, symmetry=({symmetry}))"
         else:
             # Trivial edge
             return f"(dimension={self.dimension})"
diff --git a/tests/test_edge.py b/tests/test_edge.py
index 8ff6d8dc5..aa6f457f8 100644
--- a/tests/test_edge.py
+++ b/tests/test_edge.py
@@ -29,10 +29,10 @@ def test_edge_conjugate_and_equal() -> None:
 
 def test_repr() -> None:
     a = Edge(fermion=(False, True), symmetry=(torch.tensor([False, True]), torch.tensor([0, 1])), arrow=True)
-    repr_a = "Edge(dimension=2, arrow=True, fermion=(False, True), symmetry=(tensor([False,  True]), tensor([0, 1])))"
+    repr_a = "Edge(dimension=2, arrow=True, fermion=(False,True), symmetry=([False,True],[0,1]))"
     assert repr_a == repr(a)
     b = Edge(symmetry=(torch.tensor([False, True]), torch.tensor([0, 1])))
-    repr_b = "Edge(dimension=2, symmetry=(tensor([False,  True]), tensor([0, 1])))"
+    repr_b = "Edge(dimension=2, symmetry=([False,True],[0,1]))"
     assert repr_b == repr(b)
     c = Edge(dimension=4)
     repr_c = "Edge(dimension=4)"