Save progress

3 files changed, 18 insertions, 10 deletions

diff --git a/sleekxmpp/roster.py b/sleekxmpp/roster.py
index 1972a92e..55af1e46 100644
--- a/sleekxmpp/roster.py
+++ b/sleekxmpp/roster.py
@@ -8,6 +8,8 @@
 
 import logging
 
+from sleekxmpp.xmlstream import JID
+
 
 class Roster(object):
 
@@ -64,7 +66,7 @@ class Roster(object):
         Arguments:
             key -- Return the roster for this JID.
         """
-        if not isinstance(key, str):
+        if isinstance(key, JID):
             key = key.bare
         if key not in self._rosters:
             self.add(key)
@@ -87,7 +89,7 @@ class Roster(object):
         Arguments:
             node -- The JID for the new roster node.
         """
-        if not isinstance(node, str):
+        if isinstance(node, JID):
             node = node.bare
         if node not in self._rosters:
             self._rosters[node] = RosterNode(self.xmpp, node, self.db)
@@ -169,7 +171,7 @@ class RosterNode(object):
 
         A new item entry will be created if one does not already exist.
         """
-        if not isinstance(key, str):
+        if isinstance(key, JID):
             key = key.bare
         if key not in self._jids:
             self.add(key, save=True)
@@ -228,7 +230,7 @@ class RosterNode(object):
                            if one is used.
                            Defaults to False.
         """
-        if not isinstance(jid, str):
+        if isinstance(jid, JID):
             key = jid.bare
         state = {'name': name,
                  'groups': groups or [],
diff --git a/sleekxmpp/test/sleektest.py b/sleekxmpp/test/sleektest.py
index e478e3a7..92f467fc 100644
--- a/sleekxmpp/test/sleektest.py
+++ b/sleekxmpp/test/sleektest.py
@@ -16,7 +16,8 @@ import sleekxmpp
 from sleekxmpp import ClientXMPP, ComponentXMPP
 from sleekxmpp.stanza import Message, Iq, Presence
 from sleekxmpp.test import TestSocket, TestLiveSocket
-from sleekxmpp.xmlstream import StanzaBase, ET, register_stanza_plugin
+from sleekxmpp.xmlstream import ET, register_stanza_plugin
+from sleekxmpp.xmlstream import ElementBase, StanzaBase
 from sleekxmpp.xmlstream.tostring import tostring
 from sleekxmpp.xmlstream.matcher import StanzaPath, MatcherId
 from sleekxmpp.xmlstream.matcher import MatchXMLMask, MatchXPath
@@ -227,7 +228,7 @@ class SleekTest(unittest.TestCase):
                     "Stanza:\n%s" % str(stanza))
         else:
             stanza_class = stanza.__class__
-            if isinstance(criteria, str):
+            if not isinstance(criteria, ElementBase):
                 xml = self.parse_xml(criteria)
             else:
                 xml = criteria.xml
@@ -632,7 +633,7 @@ class SleekTest(unittest.TestCase):
             self.fail("Stanza data was sent: %s" % sent)
         if sent is None:
             self.fail("No stanza was sent.")
-        
+
         xml = self.parse_xml(sent)
         self.fix_namespaces(xml, 'jabber:client')
         sent = self.xmpp._build_stanza(xml, 'jabber:client')
diff --git a/sleekxmpp/xmlstream/jid.py b/sleekxmpp/xmlstream/jid.py
index 5019a25e..36b33056 100644
--- a/sleekxmpp/xmlstream/jid.py
+++ b/sleekxmpp/xmlstream/jid.py
@@ -6,6 +6,8 @@
     See the file LICENSE for copying permission.
 """
 
+from __future__ import unicode_literals
+
 
 class JID(object):
     """
@@ -42,7 +44,9 @@ class JID(object):
         Arguments:
             jid - The new JID value.
         """
-        self._full = self._jid = str(jid)
+        if isinstance(jid, JID):
+            jid = jid.full
+        self._full = self._jid = jid
         self._domain = None
         self._resource = None
         self._user = None
@@ -123,10 +127,11 @@ class JID(object):
         return self.full
 
     def __repr__(self):
-        return str(self)
+        return self.full
 
     def __eq__(self, other):
         """
         Two JIDs are considered equal if they have the same full JID value.
         """
-        return str(other) == str(self)
+        other = JID(other)
+        return self.full == other.full