Skip to content

Commit

Permalink
Fix various typings
Browse files Browse the repository at this point in the history
  • Loading branch information
jleclanche committed Dec 14, 2022
1 parent 4530085 commit 667ed9b
Show file tree
Hide file tree
Showing 3 changed files with 27 additions and 7 deletions.
7 changes: 7 additions & 0 deletions postgres_composite_types/caster.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,10 @@
from typing import TYPE_CHECKING, Type

from psycopg2.extras import CompositeCaster

if TYPE_CHECKING:
from .composite_type import CompositeType

__all__ = ["BaseCaster"]


Expand All @@ -9,5 +14,7 @@ class BaseCaster(CompositeCaster):
instance.
"""

_composite_type_model: Type["CompositeType"]

def make(self, values):
return self._composite_type_model(*values)
20 changes: 13 additions & 7 deletions postgres_composite_types/composite_type.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import logging
from typing import Type

from django.db import connections, models
from django.db.backends.signals import connection_created
Expand Down Expand Up @@ -66,10 +67,10 @@ def __new__(cls, name, bases, attrs):
# Use an EmptyManager for everything as types cannot be queried.
meta_obj.default_manager_name = "objects"
meta_obj.base_manager_name = "objects"
attrs["objects"] = EmptyManager(model=None)
attrs["objects"] = EmptyManager(model=None) # type: ignore

ret = super().__new__(cls, name, bases, attrs)
ret.Field._composite_type_model = ret
ret.Field._composite_type_model = ret # type: ignore
return ret

def __init__(cls, name, bases, attrs):
Expand Down Expand Up @@ -147,16 +148,18 @@ class CompositeType(metaclass=CompositeTypeMeta):

# The database connection this type is registered with
registered_connection = None
_meta: Type

def __init__(self, *args, **kwargs):
if args and kwargs:
raise RuntimeError("Specify either args or kwargs but not both.")

for field in self._meta.fields:
fields = self.get_fields()
for field in fields:
setattr(self, field.name, None)

# Unpack any args as if they came from the type
for field, arg in zip(self._meta.fields, args):
for field, arg in zip(fields, args):
setattr(self, field.name, arg)

for name, value in kwargs.items():
Expand All @@ -169,21 +172,21 @@ def __repr__(self):
def __to_tuple__(self):
return tuple(
field.get_prep_value(getattr(self, field.name))
for field in self._meta.fields
for field in self.get_fields()
)

def __to_dict__(self):
return {
field.name: field.get_prep_value(getattr(self, field.name))
for field in self._meta.fields
for field in self.get_fields()
}

def __eq__(self, other):
if not isinstance(other, CompositeType):
return False
if self._meta.model != other._meta.model:
return False
for field in self._meta.fields:
for field in self.get_fields():
if getattr(self, field.name) != getattr(other, field.name):
return False
return True
Expand Down Expand Up @@ -235,3 +238,6 @@ def _get_next_or_previous_by_FIELD(self):
@classmethod
def check(cls, **kwargs):
return []

def get_fields(self):
return self._meta.fields
7 changes: 7 additions & 0 deletions postgres_composite_types/fields.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,16 @@
import json
from typing import TYPE_CHECKING, Type

from django.core.exceptions import ValidationError
from django.db.backends.postgresql.base import (
DatabaseWrapper as PostgresDatabaseWrapper,
)
from django.db.models import Field

if TYPE_CHECKING:
from .composite_type import CompositeType


__all__ = ["BaseField"]


Expand All @@ -22,6 +27,8 @@ class DummyField(Field):
class BaseField(Field):
"""Base class for the field that relates to this type."""

_composite_type_model: Type["CompositeType"]

default_error_messages = {
"bad_json": "to_python() received a string that was not valid JSON",
}
Expand Down

0 comments on commit 667ed9b

Please sign in to comment.