Implement IMnemonic.collect to support collecting a mnemonic word

Author: pjkundert

hdwallet/mnemonics/imnemonic.py

@@ -11,7 +11,7 @@
     abc
 )
 from typing import (
-    Any, Callable, Dict, Generator, List, Mapping, MutableMapping, Optional, Sequence, Set, Tuple, Union
+    Any, Callable, Collection, Dict, Generator, List, Mapping, MutableMapping, Optional, Sequence, Set, Tuple, Union
 )
 
 import os
@@ -163,6 +163,25 @@ def scan(
                 for suffix, found in self.scan( current=child, depth=max(0, depth-1), predicate=predicate ):
                     yield prefix + char + suffix, found
 
+    def options(
+        self,
+        prefix: str = '',
+        current: Optional[TrieNode] = None,
+    ) -> Generator[Tuple[bool, Set[str]], str, None]:
+        """With each symbol provided, yields the next available symbol options.
+
+        Doesn't advance unless a truthy symbol is provided via <generator>send(symbol).
+
+        Completes when the provided symbol doesn't match one of the available options.
+        """
+        last: str = ''
+        *_, (terminal, _, current) = self.find(prefix, current=current) 
+        while current is not None:
+            terminal = current.value is not current.EMPTY
+            symbol: str = yield (terminal, set(current.children))
+            if symbol:
+                current = current.children.get(symbol)
+
     def dump_lines(
         self,
         current: Optional[TrieNode] = None,
@@ -343,6 +362,9 @@ def unique( current ):
                 # Only abbreviations (not terminal words) that led to a unique terminal word
                 yield abbrev
 
+    def options(self, *args, **kwargs):
+        return self._trie.options(*args, **kwargs)
+
     def __str__(self):
         return str(self._trie)
 
@@ -663,6 +685,46 @@ def find_language(
 
         return language_indices[candidate], candidate
 
+    @classmethod
+    def collect(
+        cls,
+        languages: Optional[Collection[str]] = None,
+        wordlist_path: Optional[Dict[str, Union[str, List[str]]]] = None,
+    ) -> Generator[Tuple[Set[str], bool, Set[str]], str, None]:
+        """A generator taking input symbols, and producing a sequence of sets of possible next
+        characters in all remaining languages.
+
+        With each symbol provided, yields the remaining candidate languages, whether the symbol
+        indicated a terminal word in some language, and the available next symbols in all remaining
+        languages.
+
+        """
+        candidates: Dict[str, WordIndices] = dict(
+            (candidate, words_indices)
+            for candidate, _, words_indices in cls.wordlist_indices( wordlist_path=wordlist_path )
+            if languages is None or candidate in languages
+        )
+
+        word: str = ''
+        updaters = {
+            candidate: words_indices.options()
+            for candidate, words_indices in candidates.items()
+        }
+
+        symbol = None
+        complete = set()
+        while complete < set(updaters):
+            terminal = False
+            possible = set()
+            for candidate, updater in updaters.items():
+                try:
+                    done, available = updater.send(symbol)
+                except StopIteration:
+                    complete.add( candidate )
+                terminal |= done
+                possible |= available
+            symbol = yield (set(updaters) - complete, terminal, possible)
+
     @classmethod
     def is_valid(cls, mnemonic: Union[str, List[str]], language: Optional[str] = None, **kwargs) -> bool:
         """Checks if the given mnemonic is valid.

tests/hdwallet/mnemonics/test_mnemonics_slip39.py

@@ -2,8 +2,11 @@
 import pytest
 
 from hdwallet.exceptions import MnemonicError
+from hdwallet.mnemonics.imnemonic import (
+     Trie, WordIndices,
+)
 from hdwallet.mnemonics.slip39.mnemonic import (
-    SLIP39Mnemonic, language_parser, group_parser
+    SLIP39Mnemonic, language_parser, group_parser,
 )
 
 import shamir_mnemonic
@@ -70,7 +73,6 @@ def test_slip39_language():
         },
     }
 
-
 def test_slip39_mnemonics():
 
     # Ensure our prefix and whitespace handling works correctly

tests/test_bip39_cross_language.py

@@ -135,7 +135,7 @@ def dual_language_N_word_mnemonics(self, words=12, expected_rate=1/16, total_att
             except ChecksumError as exc:
                 # Skip invalid mnemonics (e.g., checksum failures)
                 continue
-        
+
         success_rate = len(successful_both_languages) / total_attempts
 
         print(f"{words}-word mnemonics: {len(successful_both_languages)}/{total_attempts} successful ({success_rate:.6f})")
@@ -438,7 +438,22 @@ class CustomTrieNode(TrieNode):
                   o     u     n     t                        == 16
                   u     s     e                              == 17
             h     i     e     v     e                        == 18"""
-        
+
+
+
+        options = trie.options()
+
+        assert next( options )     == (False, set( 'a' ))
+        assert options.send( 'a' ) == (False, set( 'bcd' ))
+        assert options.send( 'd' ) == (False, set( 'dj' ))
+        assert options.send( 'd' ) == (True, set( 'ir' ))
+        assert options.send( 'i' ) == (False, set( 'c' ))
+        assert options.send( 'c' ) == (False, set( 't' ))
+        assert next(options)       == (False, set( 't' ))
+        assert options.send('')    == (False, set( 't' ))
+        assert options.send( 't' ) == (True, set())
+
+
     def test_ambiguous_languages(self):
         """Test that find_language correctly detects and raises errors for ambiguous mnemonics.
 
@@ -522,6 +537,18 @@ def test_ambiguous_languages(self):
                     raise  # Re-raise unexpected errors
 
 
+def test_bip39_collection():
+
+    languages = {'english', 'french', 'spanish', 'russian'}
+
+    collect = BIP39Mnemonic.collect(languages=languages)
+    assert collect.send(None) == (languages, False, set('abcedefghijklmnopqrstuvwxyzáéíóúабвгдежзиклмнопрстуфхцчшщэюя'))
+    assert collect.send('a') == ({'english', 'french', 'spanish'}, False, set('bcedefghijlmnpqrstuvwxyzéñ'))
+    assert collect.send('d') == ({'english', 'french', 'spanish'}, False, set('adehijmoruvé'))
+    assert collect.send('d') == ({'english'},                      True , set('ir'))
+    
+
+
 def test_bip39_korean():
     # Confirm that UTF-8 Mark handling works in other languages (particularly Korean)
     (_, korean_nfc, korean_indices), = BIP39Mnemonic.wordlist_indices(