_ipaddress.py 2.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091
  1. """Implementation of the ipaddres-based network types adaptation
  2. """
  3. # psycopg/_ipaddress.py - Ipaddres-based network types adaptation
  4. #
  5. # Copyright (C) 2016-2019 Daniele Varrazzo <daniele.varrazzo@gmail.com>
  6. # Copyright (C) 2020 The Psycopg Team
  7. #
  8. # psycopg2 is free software: you can redistribute it and/or modify it
  9. # under the terms of the GNU Lesser General Public License as published
  10. # by the Free Software Foundation, either version 3 of the License, or
  11. # (at your option) any later version.
  12. #
  13. # In addition, as a special exception, the copyright holders give
  14. # permission to link this program with the OpenSSL library (or with
  15. # modified versions of OpenSSL that use the same license as OpenSSL),
  16. # and distribute linked combinations including the two.
  17. #
  18. # You must obey the GNU Lesser General Public License in all respects for
  19. # all of the code used other than OpenSSL.
  20. #
  21. # psycopg2 is distributed in the hope that it will be useful, but WITHOUT
  22. # ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
  23. # FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public
  24. # License for more details.
  25. from psycopg2.extensions import (
  26. new_type, new_array_type, register_type, register_adapter, QuotedString)
  27. from psycopg2.compat import text_type
  28. # The module is imported on register_ipaddress
  29. ipaddress = None
  30. # The typecasters are created only once
  31. _casters = None
  32. def register_ipaddress(conn_or_curs=None):
  33. """
  34. Register conversion support between `ipaddress` objects and `network types`__.
  35. :param conn_or_curs: the scope where to register the type casters.
  36. If `!None` register them globally.
  37. After the function is called, PostgreSQL :sql:`inet` values will be
  38. converted into `~ipaddress.IPv4Interface` or `~ipaddress.IPv6Interface`
  39. objects, :sql:`cidr` values into into `~ipaddress.IPv4Network` or
  40. `~ipaddress.IPv6Network`.
  41. .. __: https://www.postgresql.org/docs/current/static/datatype-net-types.html
  42. """
  43. global ipaddress
  44. import ipaddress
  45. global _casters
  46. if _casters is None:
  47. _casters = _make_casters()
  48. for c in _casters:
  49. register_type(c, conn_or_curs)
  50. for t in [ipaddress.IPv4Interface, ipaddress.IPv6Interface,
  51. ipaddress.IPv4Network, ipaddress.IPv6Network]:
  52. register_adapter(t, adapt_ipaddress)
  53. def _make_casters():
  54. inet = new_type((869,), 'INET', cast_interface)
  55. ainet = new_array_type((1041,), 'INET[]', inet)
  56. cidr = new_type((650,), 'CIDR', cast_network)
  57. acidr = new_array_type((651,), 'CIDR[]', cidr)
  58. return [inet, ainet, cidr, acidr]
  59. def cast_interface(s, cur=None):
  60. if s is None:
  61. return None
  62. # Py2 version force the use of unicode. meh.
  63. return ipaddress.ip_interface(text_type(s))
  64. def cast_network(s, cur=None):
  65. if s is None:
  66. return None
  67. return ipaddress.ip_network(text_type(s))
  68. def adapt_ipaddress(obj):
  69. return QuotedString(str(obj))