class HotClient(object): """ A Redis client wrapper that loads Lua functions and creates client methods for calling them. """ _ATOMS_FILE_NAME = "atoms.lua" _BIT_FILE_NAME = "bit.lua" _MULTI_FILE_NAME = "multi.lua" def __init__(self, client=None, *args, **kwargs): self._client = client if not self._client: kwargs.setdefault('decode_responses', True) self._client = Redis(*args, **kwargs) self._bind_atoms() self._bind_multi() def _bind_atoms(self): with open(self._get_lua_path(self._BIT_FILE_NAME)) as f: luabit = f.read() requires_luabit = ( "number_and", "number_or", "number_xor", "number_lshift", "number_rshift" ) for name, snippet in self._split_lua_file_into_funcs( self._ATOMS_FILE_NAME): if name in requires_luabit: snippet = luabit + snippet self._bind_lua_method(name, snippet) def _bind_multi(self): for name, snippet in self._split_lua_file_into_funcs("multi.lua"): self._bind_private_lua_script(name, snippet) @staticmethod def _get_lua_path(name): """ Joins the given name with the relative path of the module. """ parts = (os.path.dirname(os.path.abspath(__file__)), "lua", name) return os.path.join(*parts) def _split_lua_file_into_funcs(self, file_name): """ Returns the name / code snippet pair for each Lua function in the file under file_name. """ with open(self._get_lua_path(file_name)) as f: for func in f.read().strip().split("function "): if func: bits = func.split("\n", 1) name = bits[0].split("(")[0].strip() snippet = bits[1].rsplit("end", 1)[0].strip() yield name, snippet def _bind_lua_method(self, name, code): """ Registers the code snippet as a Lua script, and binds the script to the client as a method that can be called with the same signature as regular client methods, eg with a single key arg. """ script = self._client.register_script(code) method = lambda key, *a, **k: script(keys=[key], args=a, **k) setattr(self, name, method) def _bind_private_lua_script(self, name, code): """ Registers the code snippet as a Lua script, and binds the script to the client as a private method (eg. some_lua_func becomes a _some_lua_func method of HotClient) that can be latter wrapped in public methods with better argument and error handling. """ script = self._client.register_script(code) setattr(self, '_' + name, script) def rank_lists_by_length(self, *keys): """ Creates a temporary ZSET with LIST keys as entries and their *LLEN* as scores. Uses ZREVRANGE .. WITHSCORES, to return keys and lengths sorted from longest to shortests. :param keys: keys of the lists you want rank :return: :rtype: Ranking :raise ValueError: :raise ValueError: when not enough keys are provided """ return Ranking(self._rank_lists_by_length, keys) def rank_sets_by_cardinality(self, *keys): """ Creates a temporary ZSET with SET keys as entries and their *CARD* as scores. Uses ZREVRANGE .. WITHSCORES, to return keys and cardinalities sorted from largest to smallest. :param keys: keys of the sets you want to rank :return: :rtype: Ranking :raise ValueError: when not enough keys are provided """ return Ranking(self._rank_sets_by_cardinality, keys) def __getattr__(self, name): if name in self.__dict__: return super(HotClient, self).__getattribute__(name) return self._client.__getattribute__(name)