Implement method chaining for .absorb

GiacomoPope · Jul 29, 2024 · 2f4c21b · 2f4c21b
1 parent 32203ca
commit 2f4c21b
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -73,8 +73,9 @@ macro_rules! impl_sponge_shaker_classes {
             }
 
             #[doc=concat!("Absorb `input_bytes` into the ", stringify!($hasher_core), " state")]
-            fn absorb(&mut self, input_bytes: &[u8]) {
-                self.hasher.update(input_bytes);
+            fn absorb<'a>(mut slf: PyRefMut<'a, Self>, input_bytes: &[u8]) -> PyRefMut<'a, Self> {
+                slf.hasher.update(input_bytes);
+                slf
             }
 
             #[doc=concat!(
@@ -114,8 +115,9 @@ macro_rules! impl_sponge_shaker_classes {
             }
 
             #[doc=concat!("Absorb `input_bytes` into the ", stringify!($hasher), " state")]
-            fn absorb(&mut self, input_bytes: &[u8]) {
-                self.hasher.update(input_bytes);
+            fn absorb<'a>(mut slf: PyRefMut<'a, Self>, input_bytes: &[u8]) -> PyRefMut<'a, Self> {
+                slf.hasher.update(input_bytes);
+                slf
             }
 
             #[doc=concat!(

diff --git a/xoflib.pyi b/xoflib.pyi
@@ -2,7 +2,7 @@ class Shaker128:
     def __init__(self, input_bytes: bytes | None = None):
         ...
 
-    def absorb(self, input_bytes: bytes) -> None:
+    def absorb(self, input_bytes: bytes) -> "Shaker128":
         ...
 
     def finalize(self) -> Sponge128:
@@ -16,7 +16,7 @@ class Shaker256:
     def __init__(self, input_bytes: bytes | None = None):
         ...
 
-    def absorb(self, input_bytes: bytes) -> None:
+    def absorb(self, input_bytes: bytes) -> "Shaker256":
         ...
 
     def finalize(self) -> Sponge128:
@@ -30,7 +30,7 @@ class TurboShaker128:
     def __init__(self, domain_sep: int, input_bytes: bytes | None = None):
         ...
 
-    def absorb(self, input_bytes: bytes) -> None:
+    def absorb(self, input_bytes: bytes) -> "TurboShaker128":
         ...
 
     def finalize(self) -> Sponge128:
@@ -44,7 +44,7 @@ class TurboShaker256:
     def __init__(self, domain_sep: int, input_bytes: bytes | None = None):
         ...
 
-    def absorb(self, input_bytes: bytes) -> None:
+    def absorb(self, input_bytes: bytes) -> "TurboShaker256":
         ...
 
     def finalize(self) -> Sponge128: