-
Notifications
You must be signed in to change notification settings - Fork 95
Expand file tree
/
Copy pathbit.py
More file actions
40 lines (30 loc) · 1.25 KB
/
Copy pathbit.py
File metadata and controls
40 lines (30 loc) · 1.25 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
from sqlalchemy.dialects.postgresql.base import ischema_names
from sqlalchemy.types import UserDefinedType, Float
from sqlalchemy import Dialect, Operators
from typing import Any
class BIT(UserDefinedType):
cache_ok = True
def __init__(self, length: int | None = None) -> None:
super(UserDefinedType, self).__init__()
self.length = length
def get_col_spec(self, **kw: Any) -> str:
if self.length is None:
return 'BIT'
return 'BIT(%d)' % self.length
def bind_processor(self, dialect: Dialect) -> Any:
if dialect.__class__.__name__ == 'PGDialect_asyncpg':
import asyncpg
def process(value: Any) -> Any:
if isinstance(value, str):
return asyncpg.BitString(value)
return value
return process
else:
return super().bind_processor(dialect)
class comparator_factory(UserDefinedType.Comparator):
def hamming_distance(self, other: object) -> Operators:
return self.op('<~>', return_type=Float)(other)
def jaccard_distance(self, other: object) -> Operators:
return self.op('<%>', return_type=Float)(other)
# for reflection
ischema_names['bit'] = BIT # type: ignore