SQLAlchemy types for PostgreSQL inet types

iptypes.py

"""
This module defines additional "IP-Address" type decorators for use with
SQLAlchemy and PostgreSQL. These types will convert the PostgreSQL values
to/from Python :py:mod:`ipaddress` classes.

These classes can be used in model definitions::

    class MyTable(Base):
        ...
        ip = Column(PgIpAddress, ...)
        ...
"""

from ipaddress import (
    IPv4Address,
    IPv4Interface,
    IPv4Network,
    IPv6Address,
    IPv6Interface,
    IPv6Network,
    ip_address,
    ip_interface,
    ip_network
)
from typing import Optional, Union

from sqlalchemy.dialects.postgresql import CIDR, INET
from sqlalchemy.engine.interfaces import Dialect
from sqlalchemy.types import TypeDecorator

TInterface = Union[IPv4Interface, IPv6Interface]
TIpAddress = Union[IPv4Address, IPv6Address]
TNetwork = Union[IPv4Network, IPv6Network]


class PgIpInterface(TypeDecorator):  # type: ignore
    """
    A codec for :py:mod:`ipaddress` interfaces.
    """

    impl = INET

    def process_bind_param(
        self,
        value: Optional[TInterface],
        dialect: Dialect
    ) -> Optional[str]:
        return str(value) if value else None

    def process_result_value(
        self,
        value: Optional[str],
        dialect: Dialect
    ) -> Optional[TInterface]:
        return ip_interface(value) if value else None

    def process_literal_param(
        self,
        value: Optional[str],
        dialect: Dialect
    ) -> str:
        raise NotImplementedError('Not yet implemented')


class PgIpAddress(TypeDecorator):  # type: ignore
    """
    A codec for :py:mod:`ipaddress` IP addresses.
    """

    impl = INET

    def process_bind_param(
        self,
        value: Optional[TIpAddress],
        dialect: Dialect
    ) -> Optional[str]:
        return str(value) if value else None

    def process_result_value(
        self,
        value: Optional[str],
        dialect: Dialect
    ) -> Optional[TIpAddress]:
        if value is None:
            return None
        return ip_address(value)  # type: ignore

    def process_literal_param(
        self,
        value: Optional[str],
        dialect: Dialect
    ) -> str:
        raise NotImplementedError('Not yet implemented')


class PgCIDR(TypeDecorator):  # type: ignore
    """
    A codec for :py:mod:`ipaddress` IP networks.
    """

    impl = CIDR

    def process_bind_param(
        self,
        value: Union[TNetwork, str, None],
        dialect: Dialect
    ) -> Optional[str]:

        if value is None:
            return None

        if isinstance(value, str):
            value = ip_network(value)

        if not isinstance(value, (IPv4Network, IPv6Network)):
            raise ValueError('PgCIDR field values must be of type ip_network! '
                             'You gave me {!r}'.format(value))

        return str(value) if value else None

    def process_result_value(
        self,
        value: Optional[str],
        dialect: Dialect
    ) -> Optional[TNetwork]:
        if value is None:
            return None

        return ip_network(value)  # type: ignore

    def process_literal_param(
        self,
        value: Optional[str],
        dialect: Dialect
    ) -> str:
        raise NotImplementedError('Not yet implemented')