fix: update invert to work like connect

garvys-org · Aug 16, 2024 · 3001991 · 3001991
1 parent 19594b6
commit 3001991
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 2 deletions.
diff --git a/rustfst-python/rustfst/algorithms/inversion.py b/rustfst-python/rustfst/algorithms/inversion.py
@@ -7,15 +7,19 @@
 from rustfst.fst.vector_fst import VectorFst
 
 
-def invert(fst: VectorFst):
+def invert(fst: VectorFst) -> VectorFst:
     """
     Invert the transduction corresponding to an FST by exchanging the
     FST's input and output labels in-place.
 
     Args:
        fst: FST to be inverted.
+    Returns:
+       self
     """
 
     ret_code = lib.fst_invert(fst.ptr)
     err_msg = "Error during invert"
     check_ffi_error(ret_code, err_msg)
+
+    return fst
diff --git a/rustfst-python/rustfst/fst/vector_fst.py b/rustfst-python/rustfst/fst/vector_fst.py
@@ -724,6 +724,13 @@ def isomorphic(self, other: VectorFst) -> bool:
         return isomorphic(self, other)
 
     def invert(self) -> VectorFst:
+        """
+        Invert the transduction corresponding to an FST by exchanging the
+        FST's input and output labels in-place.
+
+        Returns:
+           self
+        """
         from rustfst.algorithms.inversion import invert
 
         return invert(self)

diff --git a/rustfst-python/tests/algorithms/test_inversion.py b/rustfst-python/tests/algorithms/test_inversion.py
@@ -24,7 +24,7 @@ def test_invert():
     tr1_4 = Tr(3, 5, 1.0, s3)
     fst1.add_tr(s2, tr1_4)
 
-    fst1.invert()
+    fst1 = fst1.invert()
 
     # Expected FST
     expected_fst = VectorFst()