_ipaddress.py 2.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990
  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-2021 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. # The module is imported on register_ipaddress
  28. ipaddress = None
  29. # The typecasters are created only once
  30. _casters = None
  31. def register_ipaddress(conn_or_curs=None):
  32. """
  33. Register conversion support between `ipaddress` objects and `network types`__.
  34. :param conn_or_curs: the scope where to register the type casters.
  35. If `!None` register them globally.
  36. After the function is called, PostgreSQL :sql:`inet` values will be
  37. converted into `~ipaddress.IPv4Interface` or `~ipaddress.IPv6Interface`
  38. objects, :sql:`cidr` values into into `~ipaddress.IPv4Network` or
  39. `~ipaddress.IPv6Network`.
  40. .. __: https://www.postgresql.org/docs/current/static/datatype-net-types.html
  41. """
  42. global ipaddress
  43. import ipaddress
  44. global _casters
  45. if _casters is None:
  46. _casters = _make_casters()
  47. for c in _casters:
  48. register_type(c, conn_or_curs)
  49. for t in [ipaddress.IPv4Interface, ipaddress.IPv6Interface,
  50. ipaddress.IPv4Network, ipaddress.IPv6Network]:
  51. register_adapter(t, adapt_ipaddress)
  52. def _make_casters():
  53. inet = new_type((869,), 'INET', cast_interface)
  54. ainet = new_array_type((1041,), 'INET[]', inet)
  55. cidr = new_type((650,), 'CIDR', cast_network)
  56. acidr = new_array_type((651,), 'CIDR[]', cidr)
  57. return [inet, ainet, cidr, acidr]
  58. def cast_interface(s, cur=None):
  59. if s is None:
  60. return None
  61. # Py2 version force the use of unicode. meh.
  62. return ipaddress.ip_interface(str(s))
  63. def cast_network(s, cur=None):
  64. if s is None:
  65. return None
  66. return ipaddress.ip_network(str(s))
  67. def adapt_ipaddress(obj):
  68. return QuotedString(str(obj))