Simple Trie Python Implementation. Comment structure based off SKLearn.

trie.py

class Trie:
    """
    Trie data structure implemented using dictionary
    
    Parameters
    ----------
    dictioanry : array
    An array of all the strings to be included in the trie
    
    Attributes
    ----------
    trie : dictionary
    dictionary with nested dictionaries representing branches of trie. Presence of
    key: '' with value '' represents termination node.
    
    Examples
    --------
    Dictionary containing cat, car, and cone
    Trie View:
        c
       / \
      a   o
     / \   \
    t   r   n
             \
              e
    Dict View: {'c': {'a': {'t': {'': ''}, 'r': {'': ''}}}, 'o': {'n': {'e': {'': ''}}}}

    """
    def __init__(self, dictionary):
        self.trie = {}

        for word in dictionary:
            currentLevel = self.trie
            for letter in word:
                currentLevel = currentLevel.setdefault(letter, {})
            currentLevel[''] = ''
           
    def search(self, word, searchType='Exists'):
        """
        Search trie for given word
        
        Parameters
        ----------
        word : string
        word to be searched for in trie
        
        searchType : 'Exists', 'Complete', 'Substring', default: 'Exists'
        Type of search to run.
        If 'Exists', will return true if search word exists as 
        full word or beginning of word in trie.
        If 'Substring', will return true if search word exists strictly 
        as beginning of word in trie.
        If 'Complete', will return true if search word exists as
        full wod in trie.
        
        Returns
        -------
        True / False
        
        """
        
        currentLevel = self.trie
        for letter in word:
            if letter in currentLevel:
                currentLevel = currentLevel[letter]
            else:
                return False
        if searchType == 'Substring':
            return '' not in currentLevel
        if searchType == 'Complete:
            return '' in currentLevel
        return True

    def getDictKeys(self, word):
        """
        Search trie for given word
        
        Parameters
        ----------
        word : string
        
        Returns
        -------
        nextLetters : array
        Array of all the potential following letters given word.
        Includes '' if complete word exists in dictionary
        
        """
        
        currentLevel = self.trie
        if not self.search(word, exists=True):
            return {}
        else:
            for letter in word:
                currentLevel = currentLevel[letter]
        nextLetters = currentLevel.keys()
        return nextLetters