How to use the asttokens.util.walk function in asttokens

i
                for i, j in enumerate(code) if re.search('REPREX_RUNNING', j)
            ]
            if x:
                code = code[(x[0] + 1):]
        scode = '\n'.join(code)
        tokes = asttokens.ASTTokens(scode, parse=True)

        def _get_one_mod(node):
            tnode = type(node).__name__
            if tnode == 'Import':
                return [i.name for i in node.names]
            elif tnode == 'ImportFrom':
                return [node.module]

        mlist = [_get_one_mod(i) for i in asttokens.util.walk(tokes.tree)]
        return {j for i in mlist if i is not None for j in i}