Small changes in connection

[diaspy.git] / diaspy / connection.py

diff --git a/diaspy/connection.py b/diaspy/connection.py
index 4f70c2f1754215d55df0bab8fa4276aa9917afd9..8654b068ede0c4c2acc4e7884fbe8ec3c404592e 100644 (file)
--- a/diaspy/connection.py
+++ b/diaspy/connection.py
@@ -19,13 +19,12 @@ class TokenError(Exception):
 
 
 class Connection():
-    """Object representing connection with the server.
-    It is pushed around internally and is considered private.
+    """Object representing connection with the pod.
     """
     _token_regex = re.compile(r'content="(.*?)"\s+name="csrf-token')
     _userinfo_regex = re.compile(r'window.current_user_attributes = ({.*})')
 
-    def __init__(self, pod, username='', password=''):
+    def __init__(self, pod, username='', password='', schema='https'):
         """
         :param pod: The complete url of the diaspora pod to use.
         :type pod: str
@@ -38,10 +37,13 @@ class Connection():
         self.session = requests.Session()
         self.login_data = {}
         self.token = ''
-        try:
-            self._setlogin(username, password)
-        except Exception as e:
-            raise LoginError('cannot create login data (caused by: {0}'.format(e))
+        try: self._setlogin(username, password)
+        except requests.exceptions.MissingSchema:
+            self.pod = '{0}://{1}'.format(schema, self.pod)
+            warnings.warn('schema was missing')
+        finally: pass
+        try: self._setlogin(username, password)
+        except Exception as e: raise LoginError('cannot create login data (caused by: {0})'.format(e))
 
     def __repr__(self):
         """Returns token string.
@@ -154,7 +156,10 @@ class Connection():
         :returns: dict -- json formatted user info.
         """
         request = self.get('bookmarklet')
-        userdata = json.loads(self._userinfo_regex.search(request.text).group(1))
+        try:
+            userdata = json.loads(self._userinfo_regex.search(request.text).group(1))
+        except AttributeError:
+            raise errors.DiaspyError('cannot find user data')
         return userdata
 
     def _fetchtoken(self):
@@ -164,9 +169,10 @@ class Connection():
         """
         request = self.get('stream')
         token = self._token_regex.search(request.text).group(1)
+        self.token = token
         return token
 
-    def get_token(self, new=False):
+    def get_token(self, fetch=False):
         """This function returns a token needed for authentication in most cases.
         Each time it is run a _fetchtoken() is called and refreshed token is stored.
 
@@ -177,8 +183,8 @@ class Connection():
         :returns: string -- token used to authenticate
         """
         try:
-            if new: self.token = self._fetchtoken()
-            if not self.token: self.token = self._fetchtoken()
+            if fetch: self._fetchtoken()
+            if not self.token: self._fetchtoken()
         except requests.exceptions.ConnectionError as e:
             warnings.warn('{0} was cought: reusing old token'.format(e))
         finally: