Merge pull request #98 from encukou/python3-parametrized-exceptions

Python3 parametrized exceptions

3 files changed, 163 insertions, 46 deletions

diff --git a/dns/exception.py b/dns/exception.py
index 60b4912..cbcdb57 100644
--- a/dns/exception.py
+++ b/dns/exception.py
@@ -15,14 +15,86 @@
 
 """Common DNS Exceptions."""
 
+
 class DNSException(Exception):
-    """Abstract base class shared by all dnspython exceptions."""
-    def __init__(self, *args):
+    """Abstract base class shared by all dnspython exceptions.
+
+    It supports two basic modes of operation:
+
+    a) Old/compatible mode is used if __init__ was called with
+       empty **kwargs.
+    In compatible mode all *args are passed to standard Python Exception class
+    as before and all *args are printed by standard __str__ implementation.
+    Class variable msg (or doc string if msg is None) is returned from str()
+    if *args is empty.
+
+    b) New/parametrized mode is used if __init__ was called with
+       non-empty **kwargs.
+    In the new mode *args has to be empty and all kwargs has to exactly match
+    set in class variable self.supp_kwargs. All kwargs are stored inside
+    self.kwargs and used in new __str__ implementation to construct
+    formated message based on self.fmt string.
+
+    In the simplest case it is enough to override supp_kwargs and fmt
+    class variables to get nice parametrized messages.
+    """
+    msg = None  # non-parametrized message
+    supp_kwargs = set()  # accepted parameters for _fmt_kwargs (sanity check)
+    fmt = None  # message parametrized with results from _fmt_kwargs
+
+    def __init__(self, *args, **kwargs):
+        self._check_params(*args, **kwargs)
+        self._check_kwargs(**kwargs)
+        self.kwargs = kwargs
+        if self.msg is None:
+            # doc string is better implicit message than empty string
+            self.msg = self.__doc__
         if args:
             super(DNSException, self).__init__(*args)
         else:
-            # doc string is better implicit message than empty string
-            super(DNSException, self).__init__(self.__doc__)
+            super(DNSException, self).__init__(self.msg)
+
+    def _check_params(self, *args, **kwargs):
+        """Old exceptions supported only args and not kwargs.
+
+        For sanity we do not allow to mix old and new behavior."""
+        if args or kwargs:
+            assert bool(args) != bool(kwargs), \
+                'keyword arguments are mutually exclusive with positional args'
+
+    def _check_kwargs(self, **kwargs):
+        if kwargs:
+            assert set(kwargs.keys()) == self.supp_kwargs, \
+                'following set of keyword args is required: %s' % (
+                    self.supp_kwargs)
+
+    def _fmt_kwargs(self, **kwargs):
+        """Format kwargs before printing them.
+
+        Resulting dictionary has to have keys necessary for str.format call
+        on fmt class variable.
+        """
+        fmtargs = {}
+        for kw, data in kwargs.items():
+            if isinstance(data, (list, set)):
+                # convert list of <someobj> to list of str(<someobj>)
+                fmtargs[kw] = list(map(str, data))
+                if len(fmtargs[kw]) == 1:
+                    # remove list brackets [] from single-item lists
+                    fmtargs[kw] = fmtargs[kw].pop()
+            else:
+                fmtargs[kw] = data
+        return fmtargs
+
+    def __str__(self):
+        if self.kwargs and self.fmt:
+            # provide custom message constructed from keyword arguments
+            fmtargs = self._fmt_kwargs(**self.kwargs)
+            return self.fmt.format(**fmtargs)
+        else:
+            # print *args directly in the same way as old DNSException
+            return super(DNSException, self).__str__()
+
 
 class FormError(DNSException):
     """DNS message is malformed."""
@@ -38,3 +110,5 @@ class TooBig(DNSException):
 
 class Timeout(DNSException):
     """The DNS operation timed out."""
+    supp_kwargs = set(['timeout'])
+    fmt = "%s after {timeout} seconds" % __doc__[:-1]
diff --git a/dns/resolver.py b/dns/resolver.py
index 387bd50..443e1d0 100644
--- a/dns/resolver.py
+++ b/dns/resolver.py
@@ -48,6 +48,21 @@ if sys.platform == 'win32':
 
 class NXDOMAIN(dns.exception.DNSException):
     """The DNS query name does not exist."""
+    supp_kwargs = set(['qname'])
+
+    def __str__(self):
+        if not 'qname' in self.kwargs:
+            return super(NXDOMAIN, self).__str__()
+
+        qname = self.kwargs['qname']
+        msg = self.__doc__[:-1]
+        if isinstance(qname, (list, set)):
+            if len(qname) > 1:
+                msg = 'None of DNS query names exist'
+                qname = list(map(str, qname))
+            else:
+                qname = qname[0]
+        return "%s: %s" % (msg, (str(qname)))
 
 class YXDOMAIN(dns.exception.DNSException):
     """The DNS query name is too long after DNAME substitution."""
@@ -58,40 +73,36 @@ class YXDOMAIN(dns.exception.DNSException):
 
 Timeout = dns.exception.Timeout
 
+
 class NoAnswer(dns.exception.DNSException):
     """The DNS response does not contain an answer to the question."""
-    def __init__(self, question=None):
-        super(dns.exception.DNSException, self).__init__()
-        self.question = question
-
-    def __str__(self):
-        message = self.__doc__
-        if self.question:
-            message = message[0:-1]
-            for q in self.question:
-                message += ' %s' % q
-        return message
+    fmt = '%s: {query}' % __doc__[:-1]
+    supp_kwargs = set(['response'])
 
+    def _fmt_kwargs(self, **kwargs):
+        return super(NoAnswer, self)._fmt_kwargs(
+            query=kwargs['response'].question)
 
 class NoNameservers(dns.exception.DNSException):
-    """No non-broken nameservers are available to answer the query."""
-    def __init__(self, errors=[]):
-        """Optionally construct message with list of servers and errors.
+    """All nameservers failed to answer the query.
 
-        @param errors: list of servers and respective errors
-        @type errors: [(server ip address, any object convertible to string)]
-        """
-        super(dns.exception.DNSException, self).__init__()
-        self.errors = errors
+    @param errors: list of servers and respective errors
+    @type errors: [(server ip address, any object convertible to string)]
+    Non-empty errors list will add explanatory message ()
+    """
+
+    msg = "All nameservers failed to answer the query."
+    fmt = "%s {query}: {errors}" % msg[:-1]
+    supp_kwargs = set(['request', 'errors'])
+
+    def _fmt_kwargs(self, **kwargs):
+        srv_msgs = []
+        for err in kwargs['errors']:
+            srv_msgs.append('Server %s %s port %s anwered %s' % (err[0],
+                            'TCP' if err[1] else 'UDP', err[2], err[3]))
+        return super(NoNameservers, self)._fmt_kwargs(
+            query=kwargs['request'].question, errors='; '.join(srv_msgs))
 
-    def __str__(self):
-        message = self.__doc__
-        if self.errors:
-            srv_msgs = []
-            for err in self.errors:
-                srv_msgs.append('Server %s %s' % (err[0], err[1]))
-            message += ' %s' % '; '.join(srv_msgs)
-        return message
 
 class NotAbsolute(dns.exception.DNSException):
     """An absolute domain name is required but a relative name was provided."""
@@ -164,9 +175,9 @@ class Answer(object):
                         if raise_on_no_answer:
                             raise NoAnswer(question=response.question)
                 if raise_on_no_answer:
-                    raise NoAnswer(question=response.question)
+                    raise NoAnswer(response=response)
         if rrset is None and raise_on_no_answer:
-            raise NoAnswer(question=response.question)
+            raise NoAnswer(response=response)
         self.canonical_name = qname
         self.rrset = rrset
         if rrset is None:
@@ -742,18 +753,18 @@ class Resolver(object):
 
     def _compute_timeout(self, start):
         now = time.time()
-        if now < start:
-            if start - now > 1:
+        duration = now - start
+        if duration < 0:
+            if duration < -1:
                 # Time going backwards is bad.  Just give up.
-                raise Timeout
+                raise Timeout(timeout=duration)
             else:
                 # Time went backwards, but only a little.  This can
                 # happen, e.g. under vmware with older linux kernels.
                 # Pretend it didn't happen.
                 now = start
-        duration = now - start
         if duration >= self.lifetime:
-            raise Timeout
+            raise Timeout(timeout=duration)
         return min(self.lifetime - duration, self.timeout)
 
     def query(self, qname, rdtype=dns.rdatatype.A, rdclass=dns.rdataclass.IN,
@@ -839,10 +850,11 @@ class Resolver(object):
             backoff = 0.10
             while response is None:
                 if len(nameservers) == 0:
-                    raise NoNameservers(errors)
+                    raise NoNameservers(request=request, errors=errors)
                 for nameserver in nameservers[:]:
                     timeout = self._compute_timeout(start)
                     try:
+                        tcp_attempt = tcp
                         if tcp:
                             response = dns.query.tcp(request, nameserver,
                                                      timeout, self.port,
@@ -855,6 +867,7 @@ class Resolver(object):
                                                      source_port=source_port)
                             if response.flags & dns.flags.TC:
                                 # Response truncated; retry with TCP.
+                                tcp_attempt = True
                                 timeout = self._compute_timeout(start)
                                 response = dns.query.tcp(request, nameserver,
                                                        timeout, self.port,
@@ -865,15 +878,17 @@ class Resolver(object):
                         # Communication failure or timeout.  Go to the
                         # next server
                         #
+                        errors.append((nameserver, tcp_attempt, self.port, ex,
+                                       response))
                         response = None
-                        errors.append((nameserver, ex))
                         continue
                     except dns.query.UnexpectedSource as ex:
                         #
                         # Who knows?  Keep going.
                         #
+                        errors.append((nameserver, tcp_attempt, self.port, ex,
+                                       response))
                         response = None
-                        errors.append((nameserver, ex))
                         continue
                     except dns.exception.FormError as ex:
                         #
@@ -882,8 +897,9 @@ class Resolver(object):
                         # continue.
                         #
                         nameservers.remove(nameserver)
+                        errors.append((nameserver, tcp_attempt, self.port, ex,
+                                       response))
                         response = None
-                        errors.append((nameserver, ex))
                         continue
                     except EOFError as ex:
                         #
@@ -893,13 +909,15 @@ class Resolver(object):
                         # mix and continue.
                         #
                         nameservers.remove(nameserver)
+                        errors.append((nameserver, tcp_attempt, self.port, ex,
+                                       response))
                         response = None
-                        errors.append((nameserver, ex))
                         continue
                     rcode = response.rcode()
                     if rcode == dns.rcode.YXDOMAIN:
                         ex = YXDOMAIN()
-                        errors.append((nameserver, ex))
+                        errors.append((nameserver, tcp_attempt, self.port, ex,
+                                       response))
                         raise ex
                     if rcode == dns.rcode.NOERROR or \
                            rcode == dns.rcode.NXDOMAIN:
@@ -911,7 +929,8 @@ class Resolver(object):
                     #
                     if rcode != dns.rcode.SERVFAIL or not self.retry_servfail:
                         nameservers.remove(nameserver)
-                    errors.append((nameserver, dns.rcode.to_text(rcode)))
+                    errors.append((nameserver, tcp_attempt, self.port,
+                                   dns.rcode.to_text(rcode), response))
                     response = None
                 if not response is None:
                     break
@@ -932,7 +951,7 @@ class Resolver(object):
             all_nxdomain = False
             break
         if all_nxdomain:
-            raise NXDOMAIN
+            raise NXDOMAIN(qname=qnames_to_try)
         answer = Answer(qname, rdtype, rdclass, response,
                         raise_on_no_answer)
         if self.cache:
diff --git a/tests/test_exceptions.py b/tests/test_exceptions.py
index 3fd7331..90a6af4 100644
--- a/tests/test_exceptions.py
+++ b/tests/test_exceptions.py
@@ -18,6 +18,12 @@ import unittest
 
 from dns.exception import DNSException
 
+
+class FormatedError(DNSException):
+    fmt = "Custom format: {parameter}"
+    supp_kwargs = set(['parameter'])
+
+
 class ExceptionTestCase(unittest.TestCase):
 
     def test_custom_message(self):
@@ -33,6 +39,24 @@ class ExceptionTestCase(unittest.TestCase):
         except DNSException as ex:
             self.assertEqual(ex.__class__.__doc__, str(ex))
 
+    def test_formatted_error(self):
+        """Exceptions with explicit format has to respect it."""
+        params = {'parameter': 'value'}
+        try:
+            raise FormatedError(**params)
+        except FormatedError as ex:
+            msg = FormatedError.fmt.format(**params)
+            self.assertEqual(msg, str(ex))
+
+    def test_kwargs_only(self):
+        """Kwargs cannot be combined with args."""
+        with self.assertRaises(AssertionError):
+            raise FormatedError(1, a=2)
+
+    def test_kwargs_unsupported(self):
+        """Only supported kwargs are accepted."""
+        with self.assertRaises(AssertionError):
+            raise FormatedError(unsupported=2)
 
 if __name__ == '__main__':
     unittest.main()