[Plugin] Improve encoding functions & storing methods

1 files changed, 60 insertions, 19 deletions

diff --git a/module/plugins/internal/Plugin.py b/module/plugins/internal/Plugin.py
index d6eac7e58..faa362dbe 100644
--- a/module/plugins/internal/Plugin.py
+++ b/module/plugins/internal/Plugin.py
@@ -6,6 +6,7 @@ import datetime
 import inspect
 import os
 import re
+import sys
 import time
 import traceback
 import urllib
@@ -17,17 +18,38 @@ if os.name is not "nt":
     import grp
     import pwd
 
+from module.common.json_layer import json_dumps, json_loads
 from module.plugins.Plugin import Abort, Fail, Reconnect, Retry, SkipDownload as Skip  #@TODO: Remove in 0.4.10
-from module.utils import fs_encode, fs_decode, html_unescape, parseFileSize as parse_size, save_join as fs_join
+from module.utils import (fs_encode, fs_decode, get_console_encoding, html_unescape,
+                          parseFileSize as parse_size, save_join as fs_join)
 
 
 #@TODO: Move to utils in 0.4.10
-def decode(string, encoding='utf8'):
-    """ Decode string to unicode with utf8 """
+def isiterable(obj):
+    return hasattr(obj, "__iter__")
+
+
+#@TODO: Move to utils in 0.4.10
+def decode(string, encoding=None):
+    """Encoded string (default to UTF-8) -> unicode string"""
     if type(string) is str:
-        return string.decode(encoding, "replace")
+        try:
+            res = unicode(string, encoding or "utf-8")
+
+        except UnicodeDecodeError, e:
+            if encoding:
+                raise UnicodeDecodeError(e)
+
+            encoding = get_console_encoding(sys.stdout.encoding)
+            res = unicode(string, encoding)
+
+    elif type(string) is unicode:
+        res = string
+
     else:
-        return unicode(string)
+        res = unicode(string)
+
+    return res
 
 
 #@TODO: Remove in 0.4.10
@@ -36,12 +58,18 @@ def _decode(*args, **kwargs):
 
 
 #@TODO: Move to utils in 0.4.10
-def encode(string, encoding='utf8'):
-    """ Decode string to utf8 """
+def encode(string, encoding=None, decoding=None):
+    """Unicode or decoded string -> encoded string (default to UTF-8)"""
     if type(string) is unicode:
-        return string.encode(encoding, "replace")
+        res = string.encode(encoding or "utf-8")
+
+    elif type(string) is str:
+        res = encode(decode(string, decoding), encoding)
+
     else:
-        return str(string)
+        res = str(string)
+
+    return res
 
 
 #@TODO: Move to utils in 0.4.10
@@ -234,7 +262,7 @@ def chunks(iterable, size):
 class Plugin(object):
     __name__    = "Plugin"
     __type__    = "plugin"
-    __version__ = "0.52"
+    __version__ = "0.53"
     __status__  = "testing"
 
     __pattern__ = r'^unmatchable$'
@@ -301,7 +329,7 @@ class Plugin(object):
 
     def log_error(self, *args, **kwargs):
         self._log("error", self.__type__, self.__name__, args)
-        if kwargs.get('trace'):
+        if self.pyload.debug and kwargs.get('trace', True):
             print "Traceback (most recent call last):"
             traceback.print_stack()
 
@@ -374,14 +402,26 @@ class Plugin(object):
         """
         Saves a value persistently to the database
         """
-        self.pyload.db.setStorage(self.__name__, key, value)
+        value = map(decode, value) if isiterable(value) else decode(value)
+        entry = json_dumps(value).encode('base64')
+        self.pyload.db.setStorage(self.__name__, key, entry)
 
 
-    def retrieve(self, key, default=None):
+    def retrieve(self, key=None, default=None):
         """
         Retrieves saved value or dict of all saved entries if key is None
         """
-        return self.pyload.db.getStorage(self.__name__, key) or default
+        entry = self.pyload.db.getStorage(self.__name__, key)
+
+        if entry:
+            if type(entry) is dict:
+                value = dict((k, json_loads(v.decode('base64'))) for k, v in value.items()) or default
+            else:
+                value = json_loads(entry.decode('base64')) or default
+        else:
+            value = entry
+
+        return value
 
 
     def delete(self, key):
@@ -456,13 +496,13 @@ class Plugin(object):
             frame = inspect.currentframe()
 
             try:
-                framefile = fs_join("tmp", self.__name__, "%s_line%s.dump.html" % (frame.f_back.f_code.co_name, frame.f_back.f_lineno))
+                framefile = fs_join("tmp", self.__name__, "%s_line%s.dump.html" %
+                                    (frame.f_back.f_code.co_name, frame.f_back.f_lineno))
 
                 if not exists(os.path.join("tmp", self.__name__)):
                     os.makedirs(os.path.join("tmp", self.__name__))
 
                 with open(framefile, "wb") as f:
-
                     f.write(encode(html))
 
             except IOError, e:
@@ -487,10 +527,11 @@ class Plugin(object):
                 value = value.strip()
 
                 if key in header:
-                    if type(header[key]) is list:
-                        header[key].append(value)
+                    header_key = header.get(key)
+                    if type(header_key) is list:
+                        header_key.append(value)
                     else:
-                        header[key] = [header[key], value]
+                        header[key] = [header_key, value]
                 else:
                     header[key] = value