From 8a298edd4e442e7d3d68ef1907d6eab524f051a7 Mon Sep 17 00:00:00 2001 From: Techmunk Date: Wed, 21 Aug 2013 23:52:54 +1000 Subject: [PATCH] Implementation of Deluge downloader. --- .../core/downloaders/deluge/__init__.py | 89 ++++ couchpotato/core/downloaders/deluge/main.py | 241 ++++++++++ .../deluge/synchronousdeluge/__init__.py | 24 + .../deluge/synchronousdeluge/client.py | 135 ++++++ .../deluge/synchronousdeluge/exceptions.py | 11 + .../deluge/synchronousdeluge/protocol.py | 38 ++ .../deluge/synchronousdeluge/rencode.py | 433 ++++++++++++++++++ .../deluge/synchronousdeluge/transfer.py | 54 +++ 8 files changed, 1025 insertions(+) create mode 100644 couchpotato/core/downloaders/deluge/__init__.py create mode 100644 couchpotato/core/downloaders/deluge/main.py create mode 100644 couchpotato/core/downloaders/deluge/synchronousdeluge/__init__.py create mode 100644 couchpotato/core/downloaders/deluge/synchronousdeluge/client.py create mode 100644 couchpotato/core/downloaders/deluge/synchronousdeluge/exceptions.py create mode 100644 couchpotato/core/downloaders/deluge/synchronousdeluge/protocol.py create mode 100644 couchpotato/core/downloaders/deluge/synchronousdeluge/rencode.py create mode 100644 couchpotato/core/downloaders/deluge/synchronousdeluge/transfer.py diff --git a/couchpotato/core/downloaders/deluge/__init__.py b/couchpotato/core/downloaders/deluge/__init__.py new file mode 100644 index 00000000..4b122b38 --- /dev/null +++ b/couchpotato/core/downloaders/deluge/__init__.py @@ -0,0 +1,89 @@ +from .main import Deluge + +def start(): + return Deluge() + +config = [{ + 'name': 'deluge', + 'groups': [ + { + 'tab': 'downloaders', + 'list': 'download_providers', + 'name': 'deluge', + 'label': 'Deluge', + 'description': 'Use Deluge to download torrents.', + 'wizard': True, + 'options': [ + { + 'name': 'enabled', + 'default': 0, + 'type': 'enabler', + 'radio_group': 'torrent', + }, + { + 'name': 'host', + 'default': 'localhost:58846', + 'description': 'Hostname with port. Usually localhost:58846', + }, + { + 'name': 'username', + }, + { + 'name': 'password', + 'type': 'password', + }, + { + 'name': 'paused', + 'type': 'bool', + 'default': False, + 'description': 'Add the torrent paused.', + }, + { + 'name': 'directory', + 'type': 'directory', + 'description': 'Download to this directory. Keep empty for default Deluge download directory.', + }, + { + 'name': 'completed_directory', + 'type': 'directory', + 'description': 'Move completed torrent to this directory. Keep empty for default Deluge options.', + 'advanced': True, + }, + { + 'name': 'label', + 'description': 'Label to add to torrents in the Deluge UI.', + }, + { + 'name': 'remove_complete', + 'label': 'Remove torrent', + 'type': 'bool', + 'default': True, + 'advanced': True, + 'description': 'Remove the torrent from Deluge after it has finished seeding.', + }, + { + 'name': 'delete_files', + 'label': 'Remove files', + 'default': True, + 'type': 'bool', + 'advanced': True, + 'description': 'Also remove the leftover files.', + }, + { + 'name': 'manual', + 'default': 0, + 'type': 'bool', + 'advanced': True, + 'description': 'Disable this downloader for automated searches, but use it when I manually send a release.', + }, + { + 'name': 'delete_failed', + 'default': True, + 'advanced': True, + 'type': 'bool', + 'description': 'Delete a release after the download has failed.', + }, + ], + } + ], +}] diff --git a/couchpotato/core/downloaders/deluge/main.py b/couchpotato/core/downloaders/deluge/main.py new file mode 100644 index 00000000..a990b175 --- /dev/null +++ b/couchpotato/core/downloaders/deluge/main.py @@ -0,0 +1,241 @@ +from base64 import b64encode +from couchpotato.core.helpers.variable import tryInt, tryFloat +from couchpotato.core.downloaders.base import Downloader, StatusList +from couchpotato.core.helpers.encoding import isInt +from couchpotato.core.logger import CPLog +from couchpotato.environment import Env +from datetime import timedelta + +from synchronousdeluge import DelugeClient + +import os.path +import traceback + +log = CPLog(__name__) + +class Deluge(Downloader): + + protocol = ['torrent', 'torrent_magnet'] + log = CPLog(__name__) + drpc = None + + def connect(self): + # Load host from config and split out port. + host = self.conf('host').split(':') + if not isInt(host[1]): + log.error('Config properties are not filled in correctly, port is missing.') + return False + + if not self.drpc: + self.drpc = DelugeRPC(host[0], port = host[1], username = self.conf('username'), password = self.conf('password')) + + return self.drpc + + def download(self, data, movie, filedata = None): + + log.info('Sending "%s" (%s) to Deluge.', (data.get('name'), data.get('protocol'))) + + if not self.connect(): + return False + + if not filedata and data.get('protocol') == 'torrent': + log.error('Failed sending torrent, no data') + return False + + # Set parameters for Deluge + options = { + 'add_paused': self.conf('paused', default = 0), + 'label': self.conf('label') + } + + if self.conf('directory'): + if os.path.isdir(self.conf('directory')): + options['download_location'] = self.conf('directory') + else: + log.error('Download directory from Deluge settings: %s doesn\'t exist', self.conf('directory')) + + if self.conf('completed_directory'): + if os.path.isdir(self.conf('completed_directory')): + options['move_completed'] = 1 + options['move_completed_path'] = self.conf('completed_directory') + else: + log.error('Download directory from Deluge settings: %s doesn\'t exist', self.conf('directory')) + + if data.get('seed_ratio'): + options['stop_at_ratio'] = 1 + options['stop_ratio'] = tryFloat(data.get('seed_ratio')) + +# Deluge only has seed time as a global option. Might be added in +# in a future API release. +# if data.get('seed_time'): + + # Send request to Deluge + if data.get('protocol') == 'torrent_magnet': + remote_torrent = self.drpc.add_torrent_magnet(data.get('url'), options) + else: + remote_torrent = self.drpc.add_torrent_file(movie, b64encode(filedata), options) + + if not remote_torrent: + log.error('Failed sending torrent to Deluge') + return False + + log.info('Torrent sent to Deluge successfully.') + return self.downloadReturnId(remote_torrent) + + def getAllDownloadStatus(self): + + log.debug('Checking Deluge download status.') + + if not self.connect(): + return False + + statuses = StatusList(self) + + queue = self.drpc.get_alltorrents() + + if not (queue and queue.get('torrents')): + log.debug('Nothing in queue or error') + return False + + for torrent_id in queue: + item = queue[torrent_id] + log.debug('name=%s / id=%s / save_path=%s / hash=%s / progress=%s / state=%s / eta=%s / ratio=%s / conf_ratio=%s/ is_seed=%s / is_finished=%s', (item['name'], item['hash'], item['save_path'], item['hash'], item['progress'], item['state'], item['eta'], item['ratio'], self.conf('ratio'), item['is_seed'], item['is_finished'])) + + if not os.path.isdir(Env.setting('from', 'renamer')): + log.error('Renamer "from" folder doesn\'t to exist.') + return + + status = 'busy' + # Deluge seems to set both is_seed and is_finished once everything has been downloaded. + if item['is_seed'] or item['is_finished']: + status = 'seeding' + elif item['is_seed'] and item['is_finished'] and item['paused']: + status = 'completed' + + download_dir = item['save_path'] + if item['move_on_completed']: + download_dir = item['move_completed_path'] + + statuses.append({ + 'id': item['hash'], + 'name': item['name'], + 'status': status, + 'original_status': item['state'], + 'seed_ratio': item['ratio'], + 'timeleft': str(timedelta(seconds = item['eta'])), + 'folder': os.path.join(download_dir, item['name']), + }) + + return statuses + + def pause(self, item, pause = True): + if pause: + return self.drpc.pause_torrent([item['id']]) + else: + return self.drpc.resume_torrent([item['id']]) + + def removeFailed(self, item): + log.info('%s failed downloading, deleting...', item['name']) + return self.drpc.remove_torrent(item['id'], True) + + def processComplete(self, item, delete_files = False): + log.debug('Requesting Deluge to remove the torrent %s%s.', (item['name'], ' and cleanup the downloaded files' if delete_files else '')) + return self.drpc.remove_torrent(item['id'], remove_local_data = delete_files) + +class DelugeRPC(object): + + host = 'localhost' + port = 58846 + username = None + password = None + client = None + + def __init__(self, host = 'localhost', port = 58846, username = None, password = None): + super(DelugeRPC, self).__init__() + + self.host = host + self.port = port + self.username = username + self.password = password + + def connect(self): + self.client = DelugeClient() + self.client.connect(self.host, int(self.port), self.username, self.password) + + def add_torrent_magnet(self, torrent, options): + torrent_id = False + try: + self.connect() + torrent_id = self.client.core.add_torrent_magnet(torrent, options).get() + if options['label']: + self.client.label.set_torrent(torrent_id, options['label']).get() + except Exception, err: + log.error('Failed to add torrent magnet: %s %s', err, traceback.format_exc()) + finally: + if self.client: + self.disconnect() + + return torrent_id + + def add_torrent_file(self, movie, torrent, options): + torrent_id = False + try: + self.connect() + torrent_id = self.client.core.add_torrent_file(movie, torrent, options).get() + if options['label']: + self.client.label.set_torrent(torrent_id, options['label']).get() + except Exception, err: + log.error('Failed to add torrent file: %s %s', err, traceback.format_exc()) + finally: + if self.client: + self.disconnect() + + return torrent_id + + def get_alltorrents(self): + ret = False + try: + self.connect() + ret = self.client.core.get_torrents_status({}, {}).get() + except Exception, err: + log.error('Failed to get all torrents: %s %s', err, traceback.format_exc()) + finally: + if self.client: + self.disconnect() + return ret + + def pause_torrent(self, torrent_ids): + try: + self.connect() + self.client.core.pause_torrent(torrent_ids).get() + except Exception, err: + log.error('Failed to pause torrent: %s %s', err, traceback.format_exc()) + finally: + if self.client: + self.disconnect() + + def resume_torrent(self, torrent_ids): + try: + self.connect() + self.client.core.resume_torrent(torrent_ids).get() + except Exception, err: + log.error('Failed to resume torrent: %s %s', err, traceback.format_exc()) + finally: + if self.client: + self.disconnect() + + def remove_torrent(self, torrent_id, remove_local_data): + ret = False + try: + self.connect() + ret = self.client.core.remove_torrent(torrent_id, remove_local_data).get() + except Exception, err: + log.error('Failed to remove torrent: %s %s', err, traceback.format_exc()) + finally: + if self.client: + self.disconnect() + return ret + + def disconnect(self): + self.client.disconnect() + diff --git a/couchpotato/core/downloaders/deluge/synchronousdeluge/__init__.py b/couchpotato/core/downloaders/deluge/synchronousdeluge/__init__.py new file mode 100644 index 00000000..a6fbcdd8 --- /dev/null +++ b/couchpotato/core/downloaders/deluge/synchronousdeluge/__init__.py @@ -0,0 +1,24 @@ +"""A synchronous implementation of the Deluge RPC protocol + based on gevent-deluge by Christopher Rosell. + + https://github.com/chrippa/gevent-deluge + +Example usage: + + from synchronousdeluge import DelgueClient + + client = DelugeClient() + client.connect() + + # Wait for value + download_location = client.core.get_config_value("download_location").get() +""" + + +__title__ = "synchronous-deluge" +__version__ = "0.1" +__author__ = "Christian Dale" + +from .client import DelugeClient +from .exceptions import DelugeRPCError + diff --git a/couchpotato/core/downloaders/deluge/synchronousdeluge/client.py b/couchpotato/core/downloaders/deluge/synchronousdeluge/client.py new file mode 100644 index 00000000..363bd855 --- /dev/null +++ b/couchpotato/core/downloaders/deluge/synchronousdeluge/client.py @@ -0,0 +1,135 @@ +import os + +from collections import defaultdict +from itertools import imap + +from .exceptions import DelugeRPCError +from .protocol import DelugeRPCRequest, DelugeRPCResponse +from .transfer import DelugeTransfer + +__all__ = ["DelugeClient"] + + +RPC_RESPONSE = 1 +RPC_ERROR = 2 +RPC_EVENT = 3 + + +class DelugeClient(object): + def __init__(self): + """A deluge client session.""" + self.transfer = DelugeTransfer() + self.modules = [] + self._request_counter = 0 + + def _get_local_auth(self): + xdg_config = os.path.expanduser(os.environ.get("XDG_CONFIG_HOME", "~/.config")) + config_home = os.path.join(xdg_config, "deluge") + auth_file = os.path.join(config_home, "auth") + + username = password = "" + with open(auth_file) as fd: + for line in fd: + if line.startswith("#"): + continue + + auth = line.split(":") + if len(auth) >= 2 and auth[0] == "localclient": + username, password = auth[0], auth[1] + break + + return username, password + + def _create_module_method(self, module, method): + fullname = "{0}.{1}".format(module, method) + + def func(obj, *args, **kwargs): + return self.remote_call(fullname, *args, **kwargs) + + func.__name__ = method + + return func + + def _introspect(self): + self.modules = [] + + methods = self.remote_call("daemon.get_method_list").get() + methodmap = defaultdict(dict) + splitter = lambda v: v.split(".") + + for module, method in imap(splitter, methods): + methodmap[module][method] = self._create_module_method(module, method) + + for module, methods in methodmap.items(): + clsname = "DelugeModule{0}".format(module.capitalize()) + cls = type(clsname, (), methods) + setattr(self, module, cls()) + self.modules.append(module) + + def remote_call(self, method, *args, **kwargs): + req = DelugeRPCRequest(self._request_counter, method, *args, **kwargs) + message = next(self.transfer.send_request(req)) + + response = DelugeRPCResponse() + + if not isinstance(message, tuple): + return + + if len(message) < 3: + return + + message_type = message[0] + +# if message_type == RPC_EVENT: +# event = message[1] +# values = message[2] +# +# if event in self._event_handlers: +# for handler in self._event_handlers[event]: +# gevent.spawn(handler, *values) +# +# elif message_type in (RPC_RESPONSE, RPC_ERROR): + if message_type in (RPC_RESPONSE, RPC_ERROR): + request_id = message[1] + value = message[2] + + if request_id == self._request_counter : + if message_type == RPC_RESPONSE: + response.set(value) + elif message_type == RPC_ERROR: + err = DelugeRPCError(*value) + response.set_exception(err) + + self._request_counter += 1 + return response + + def connect(self, host="127.0.0.1", port=58846, username="", password=""): + """Connects to a daemon process. + + :param host: str, the hostname of the daemon + :param port: int, the port of the daemon + :param username: str, the username to login with + :param password: str, the password to login with + """ + + # Connect transport + self.transfer.connect((host, port)) + + # Attempt to fetch local auth info if needed + if not username and host in ("127.0.0.1", "localhost"): + username, password = self._get_local_auth() + + # Authenticate + self.remote_call("daemon.login", username, password).get() + + # Introspect available methods + self._introspect() + + @property + def connected(self): + return self.transfer.connected + + def disconnect(self): + """Disconnects from the daemon.""" + self.transfer.disconnect() + diff --git a/couchpotato/core/downloaders/deluge/synchronousdeluge/exceptions.py b/couchpotato/core/downloaders/deluge/synchronousdeluge/exceptions.py new file mode 100644 index 00000000..da6cf022 --- /dev/null +++ b/couchpotato/core/downloaders/deluge/synchronousdeluge/exceptions.py @@ -0,0 +1,11 @@ +__all__ = ["DelugeRPCError"] + +class DelugeRPCError(Exception): + def __init__(self, name, msg, traceback): + self.name = name + self.msg = msg + self.traceback = traceback + + def __str__(self): + return "{0}: {1}: {2}".format(self.__class__.__name__, self.name, self.msg) + diff --git a/couchpotato/core/downloaders/deluge/synchronousdeluge/protocol.py b/couchpotato/core/downloaders/deluge/synchronousdeluge/protocol.py new file mode 100644 index 00000000..756d4dfc --- /dev/null +++ b/couchpotato/core/downloaders/deluge/synchronousdeluge/protocol.py @@ -0,0 +1,38 @@ +__all__ = ["DelugeRPCRequest", "DelugeRPCResponse"] + +class DelugeRPCRequest(object): + def __init__(self, request_id, method, *args, **kwargs): + self.request_id = request_id + self.method = method + self.args = args + self.kwargs = kwargs + + def format(self): + return (self.request_id, self.method, self.args, self.kwargs) + +class DelugeRPCResponse(object): + def __init__(self): + self.value = None + self._exception = None + + def successful(self): + return self._exception is None + + @property + def exception(self): + if self._exception is not None: + return self._exception + + def set(self, value=None): + self.value = value + self._exception = None + + def set_exception(self, exception): + self._exception = exception + + def get(self): + if self._exception is None: + return self.value + else: + raise self._exception + diff --git a/couchpotato/core/downloaders/deluge/synchronousdeluge/rencode.py b/couchpotato/core/downloaders/deluge/synchronousdeluge/rencode.py new file mode 100644 index 00000000..e58c7154 --- /dev/null +++ b/couchpotato/core/downloaders/deluge/synchronousdeluge/rencode.py @@ -0,0 +1,433 @@ + +""" +rencode -- Web safe object pickling/unpickling. + +Public domain, Connelly Barnes 2006-2007. + +The rencode module is a modified version of bencode from the +BitTorrent project. For complex, heterogeneous data structures with +many small elements, r-encodings take up significantly less space than +b-encodings: + + >>> len(rencode.dumps({'a':0, 'b':[1,2], 'c':99})) + 13 + >>> len(bencode.bencode({'a':0, 'b':[1,2], 'c':99})) + 26 + +The rencode format is not standardized, and may change with different +rencode module versions, so you should check that you are using the +same rencode version throughout your project. +""" + +__version__ = '1.0.1' +__all__ = ['dumps', 'loads'] + +# Original bencode module by Petru Paler, et al. +# +# Modifications by Connelly Barnes: +# +# - Added support for floats (sent as 32-bit or 64-bit in network +# order), bools, None. +# - Allowed dict keys to be of any serializable type. +# - Lists/tuples are always decoded as tuples (thus, tuples can be +# used as dict keys). +# - Embedded extra information in the 'typecodes' to save some space. +# - Added a restriction on integer length, so that malicious hosts +# cannot pass us large integers which take a long time to decode. +# +# Licensed by Bram Cohen under the "MIT license": +# +# "Copyright (C) 2001-2002 Bram Cohen +# +# Permission is hereby granted, free of charge, to any person +# obtaining a copy of this software and associated documentation files +# (the "Software"), to deal in the Software without restriction, +# including without limitation the rights to use, copy, modify, merge, +# publish, distribute, sublicense, and/or sell copies of the Software, +# and to permit persons to whom the Software is furnished to do so, +# subject to the following conditions: +# +# The above copyright notice and this permission notice shall be +# included in all copies or substantial portions of the Software. +# +# The Software is provided "AS IS", without warranty of any kind, +# express or implied, including but not limited to the warranties of +# merchantability, fitness for a particular purpose and +# noninfringement. In no event shall the authors or copyright holders +# be liable for any claim, damages or other liability, whether in an +# action of contract, tort or otherwise, arising from, out of or in +# connection with the Software or the use or other dealings in the +# Software." +# +# (The rencode module is licensed under the above license as well). +# + +import struct +import string +from threading import Lock + +# Default number of bits for serialized floats, either 32 or 64 (also a parameter for dumps()). +DEFAULT_FLOAT_BITS = 32 + +# Maximum length of integer when written as base 10 string. +MAX_INT_LENGTH = 64 + +# The bencode 'typecodes' such as i, d, etc have been extended and +# relocated on the base-256 character set. +CHR_LIST = chr(59) +CHR_DICT = chr(60) +CHR_INT = chr(61) +CHR_INT1 = chr(62) +CHR_INT2 = chr(63) +CHR_INT4 = chr(64) +CHR_INT8 = chr(65) +CHR_FLOAT32 = chr(66) +CHR_FLOAT64 = chr(44) +CHR_TRUE = chr(67) +CHR_FALSE = chr(68) +CHR_NONE = chr(69) +CHR_TERM = chr(127) + +# Positive integers with value embedded in typecode. +INT_POS_FIXED_START = 0 +INT_POS_FIXED_COUNT = 44 + +# Dictionaries with length embedded in typecode. +DICT_FIXED_START = 102 +DICT_FIXED_COUNT = 25 + +# Negative integers with value embedded in typecode. +INT_NEG_FIXED_START = 70 +INT_NEG_FIXED_COUNT = 32 + +# Strings with length embedded in typecode. +STR_FIXED_START = 128 +STR_FIXED_COUNT = 64 + +# Lists with length embedded in typecode. +LIST_FIXED_START = STR_FIXED_START+STR_FIXED_COUNT +LIST_FIXED_COUNT = 64 + +def decode_int(x, f): + f += 1 + newf = x.index(CHR_TERM, f) + if newf - f >= MAX_INT_LENGTH: + raise ValueError('overflow') + try: + n = int(x[f:newf]) + except (OverflowError, ValueError): + n = long(x[f:newf]) + if x[f] == '-': + if x[f + 1] == '0': + raise ValueError + elif x[f] == '0' and newf != f+1: + raise ValueError + return (n, newf+1) + +def decode_intb(x, f): + f += 1 + return (struct.unpack('!b', x[f:f+1])[0], f+1) + +def decode_inth(x, f): + f += 1 + return (struct.unpack('!h', x[f:f+2])[0], f+2) + +def decode_intl(x, f): + f += 1 + return (struct.unpack('!l', x[f:f+4])[0], f+4) + +def decode_intq(x, f): + f += 1 + return (struct.unpack('!q', x[f:f+8])[0], f+8) + +def decode_float32(x, f): + f += 1 + n = struct.unpack('!f', x[f:f+4])[0] + return (n, f+4) + +def decode_float64(x, f): + f += 1 + n = struct.unpack('!d', x[f:f+8])[0] + return (n, f+8) + +def decode_string(x, f): + colon = x.index(':', f) + try: + n = int(x[f:colon]) + except (OverflowError, ValueError): + n = long(x[f:colon]) + if x[f] == '0' and colon != f+1: + raise ValueError + colon += 1 + s = x[colon:colon+n] + try: + t = s.decode("utf8") + if len(t) != len(s): + s = t + except UnicodeDecodeError: + pass + return (s, colon+n) + +def decode_list(x, f): + r, f = [], f+1 + while x[f] != CHR_TERM: + v, f = decode_func[x[f]](x, f) + r.append(v) + return (tuple(r), f + 1) + +def decode_dict(x, f): + r, f = {}, f+1 + while x[f] != CHR_TERM: + k, f = decode_func[x[f]](x, f) + r[k], f = decode_func[x[f]](x, f) + return (r, f + 1) + +def decode_true(x, f): + return (True, f+1) + +def decode_false(x, f): + return (False, f+1) + +def decode_none(x, f): + return (None, f+1) + +decode_func = {} +decode_func['0'] = decode_string +decode_func['1'] = decode_string +decode_func['2'] = decode_string +decode_func['3'] = decode_string +decode_func['4'] = decode_string +decode_func['5'] = decode_string +decode_func['6'] = decode_string +decode_func['7'] = decode_string +decode_func['8'] = decode_string +decode_func['9'] = decode_string +decode_func[CHR_LIST ] = decode_list +decode_func[CHR_DICT ] = decode_dict +decode_func[CHR_INT ] = decode_int +decode_func[CHR_INT1 ] = decode_intb +decode_func[CHR_INT2 ] = decode_inth +decode_func[CHR_INT4 ] = decode_intl +decode_func[CHR_INT8 ] = decode_intq +decode_func[CHR_FLOAT32] = decode_float32 +decode_func[CHR_FLOAT64] = decode_float64 +decode_func[CHR_TRUE ] = decode_true +decode_func[CHR_FALSE ] = decode_false +decode_func[CHR_NONE ] = decode_none + +def make_fixed_length_string_decoders(): + def make_decoder(slen): + def f(x, f): + s = x[f+1:f+1+slen] + try: + t = s.decode("utf8") + if len(t) != len(s): + s = t + except UnicodeDecodeError: + pass + return (s, f+1+slen) + return f + for i in range(STR_FIXED_COUNT): + decode_func[chr(STR_FIXED_START+i)] = make_decoder(i) + +make_fixed_length_string_decoders() + +def make_fixed_length_list_decoders(): + def make_decoder(slen): + def f(x, f): + r, f = [], f+1 + for i in range(slen): + v, f = decode_func[x[f]](x, f) + r.append(v) + return (tuple(r), f) + return f + for i in range(LIST_FIXED_COUNT): + decode_func[chr(LIST_FIXED_START+i)] = make_decoder(i) + +make_fixed_length_list_decoders() + +def make_fixed_length_int_decoders(): + def make_decoder(j): + def f(x, f): + return (j, f+1) + return f + for i in range(INT_POS_FIXED_COUNT): + decode_func[chr(INT_POS_FIXED_START+i)] = make_decoder(i) + for i in range(INT_NEG_FIXED_COUNT): + decode_func[chr(INT_NEG_FIXED_START+i)] = make_decoder(-1-i) + +make_fixed_length_int_decoders() + +def make_fixed_length_dict_decoders(): + def make_decoder(slen): + def f(x, f): + r, f = {}, f+1 + for j in range(slen): + k, f = decode_func[x[f]](x, f) + r[k], f = decode_func[x[f]](x, f) + return (r, f) + return f + for i in range(DICT_FIXED_COUNT): + decode_func[chr(DICT_FIXED_START+i)] = make_decoder(i) + +make_fixed_length_dict_decoders() + +def encode_dict(x,r): + r.append(CHR_DICT) + for k, v in x.items(): + encode_func[type(k)](k, r) + encode_func[type(v)](v, r) + r.append(CHR_TERM) + + +def loads(x): + try: + r, l = decode_func[x[0]](x, 0) + except (IndexError, KeyError): + raise ValueError + if l != len(x): + raise ValueError + return r + +from types import StringType, IntType, LongType, DictType, ListType, TupleType, FloatType, NoneType, UnicodeType + +def encode_int(x, r): + if 0 <= x < INT_POS_FIXED_COUNT: + r.append(chr(INT_POS_FIXED_START+x)) + elif -INT_NEG_FIXED_COUNT <= x < 0: + r.append(chr(INT_NEG_FIXED_START-1-x)) + elif -128 <= x < 128: + r.extend((CHR_INT1, struct.pack('!b', x))) + elif -32768 <= x < 32768: + r.extend((CHR_INT2, struct.pack('!h', x))) + elif -2147483648 <= x < 2147483648: + r.extend((CHR_INT4, struct.pack('!l', x))) + elif -9223372036854775808 <= x < 9223372036854775808: + r.extend((CHR_INT8, struct.pack('!q', x))) + else: + s = str(x) + if len(s) >= MAX_INT_LENGTH: + raise ValueError('overflow') + r.extend((CHR_INT, s, CHR_TERM)) + +def encode_float32(x, r): + r.extend((CHR_FLOAT32, struct.pack('!f', x))) + +def encode_float64(x, r): + r.extend((CHR_FLOAT64, struct.pack('!d', x))) + +def encode_bool(x, r): + r.extend({False: CHR_FALSE, True: CHR_TRUE}[bool(x)]) + +def encode_none(x, r): + r.extend(CHR_NONE) + +def encode_string(x, r): + if len(x) < STR_FIXED_COUNT: + r.extend((chr(STR_FIXED_START + len(x)), x)) + else: + r.extend((str(len(x)), ':', x)) + +def encode_unicode(x, r): + encode_string(x.encode("utf8"), r) + +def encode_list(x, r): + if len(x) < LIST_FIXED_COUNT: + r.append(chr(LIST_FIXED_START + len(x))) + for i in x: + encode_func[type(i)](i, r) + else: + r.append(CHR_LIST) + for i in x: + encode_func[type(i)](i, r) + r.append(CHR_TERM) + +def encode_dict(x,r): + if len(x) < DICT_FIXED_COUNT: + r.append(chr(DICT_FIXED_START + len(x))) + for k, v in x.items(): + encode_func[type(k)](k, r) + encode_func[type(v)](v, r) + else: + r.append(CHR_DICT) + for k, v in x.items(): + encode_func[type(k)](k, r) + encode_func[type(v)](v, r) + r.append(CHR_TERM) + +encode_func = {} +encode_func[IntType] = encode_int +encode_func[LongType] = encode_int +encode_func[StringType] = encode_string +encode_func[ListType] = encode_list +encode_func[TupleType] = encode_list +encode_func[DictType] = encode_dict +encode_func[NoneType] = encode_none +encode_func[UnicodeType] = encode_unicode + +lock = Lock() + +try: + from types import BooleanType + encode_func[BooleanType] = encode_bool +except ImportError: + pass + +def dumps(x, float_bits=DEFAULT_FLOAT_BITS): + """ + Dump data structure to str. + + Here float_bits is either 32 or 64. + """ + lock.acquire() + try: + if float_bits == 32: + encode_func[FloatType] = encode_float32 + elif float_bits == 64: + encode_func[FloatType] = encode_float64 + else: + raise ValueError('Float bits (%d) is not 32 or 64' % float_bits) + r = [] + encode_func[type(x)](x, r) + finally: + lock.release() + return ''.join(r) + +def test(): + f1 = struct.unpack('!f', struct.pack('!f', 25.5))[0] + f2 = struct.unpack('!f', struct.pack('!f', 29.3))[0] + f3 = struct.unpack('!f', struct.pack('!f', -0.6))[0] + L = (({'a':15, 'bb':f1, 'ccc':f2, '':(f3,(),False,True,'')},('a',10**20),tuple(range(-100000,100000)),'b'*31,'b'*62,'b'*64,2**30,2**33,2**62,2**64,2**30,2**33,2**62,2**64,False,False, True, -1, 2, 0),) + assert loads(dumps(L)) == L + d = dict(zip(range(-100000,100000),range(-100000,100000))) + d.update({'a':20, 20:40, 40:41, f1:f2, f2:f3, f3:False, False:True, True:False}) + L = (d, {}, {5:6}, {7:7,True:8}, {9:10, 22:39, 49:50, 44: ''}) + assert loads(dumps(L)) == L + L = ('', 'a'*10, 'a'*100, 'a'*1000, 'a'*10000, 'a'*100000, 'a'*1000000, 'a'*10000000) + assert loads(dumps(L)) == L + L = tuple([dict(zip(range(n),range(n))) for n in range(100)]) + ('b',) + assert loads(dumps(L)) == L + L = tuple([dict(zip(range(n),range(-n,0))) for n in range(100)]) + ('b',) + assert loads(dumps(L)) == L + L = tuple([tuple(range(n)) for n in range(100)]) + ('b',) + assert loads(dumps(L)) == L + L = tuple(['a'*n for n in range(1000)]) + ('b',) + assert loads(dumps(L)) == L + L = tuple(['a'*n for n in range(1000)]) + (None,True,None) + assert loads(dumps(L)) == L + assert loads(dumps(None)) == None + assert loads(dumps({None:None})) == {None:None} + assert 1e-10