Skip to content

Commit c9eea04

Browse files
1pxone1pxone
committed
Improve typings
1 parent 199d5c5 commit c9eea04

File tree

7 files changed

+98
-82
lines changed

7 files changed

+98
-82
lines changed

‎src/rushdb/api/labels.py

+6-4
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import typing
12
from typing import List, Optional
23

34
from ..models.search_query import SearchQuery
@@ -14,10 +15,11 @@ def list(
1415
transaction: Optional[Transaction] = None,
1516
) -> List[str]:
1617
"""List all labels."""
17-
headers = Transaction._build_transaction_header(
18-
transaction.id if transaction else None
19-
)
18+
headers = Transaction._build_transaction_header(transaction)
2019

2120
return self.client._make_request(
22-
"POST", "/api/v1/labels", data=query or {}, headers=headers
21+
"POST",
22+
"/api/v1/labels",
23+
data=typing.cast(typing.Dict[str, typing.Any], query or {}),
24+
headers=headers,
2325
)

‎src/rushdb/api/properties.py

+9-13
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import typing
12
from typing import List, Literal, Optional
23

34
from ..models.property import Property, PropertyValuesData
@@ -15,21 +16,20 @@ def find(
1516
transaction: Optional[Transaction] = None,
1617
) -> List[Property]:
1718
"""List all properties."""
18-
headers = Transaction._build_transaction_header(
19-
transaction.id if transaction else None
20-
)
19+
headers = Transaction._build_transaction_header(transaction)
2120

2221
return self.client._make_request(
23-
"POST", "/api/v1/properties", query or {}, headers
22+
"POST",
23+
"/api/v1/properties",
24+
typing.cast(typing.Dict[str, typing.Any], query or {}),
25+
headers,
2426
)
2527

2628
def find_by_id(
2729
self, property_id: str, transaction: Optional[Transaction] = None
2830
) -> Property:
2931
"""Get a property by ID."""
30-
headers = Transaction._build_transaction_header(
31-
transaction.id if transaction else None
32-
)
32+
headers = Transaction._build_transaction_header(transaction)
3333

3434
return self.client._make_request(
3535
"GET", f"/api/v1/properties/{property_id}", headers=headers
@@ -39,9 +39,7 @@ def delete(
3939
self, property_id: str, transaction: Optional[Transaction] = None
4040
) -> None:
4141
"""Delete a property."""
42-
headers = Transaction._build_transaction_header(
43-
transaction.id if transaction else None
44-
)
42+
headers = Transaction._build_transaction_header(transaction)
4543

4644
return self.client._make_request(
4745
"DELETE", f"/api/v1/properties/{property_id}", headers=headers
@@ -56,9 +54,7 @@ def values(
5654
transaction: Optional[Transaction] = None,
5755
) -> PropertyValuesData:
5856
"""Get values data for a property."""
59-
headers = Transaction._build_transaction_header(
60-
transaction.id if transaction else None
61-
)
57+
headers = Transaction._build_transaction_header(transaction)
6258

6359
return self.client._make_request(
6460
"GET",

‎src/rushdb/api/records.py

+52-40
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import typing
12
from typing import Any, Dict, List, Optional, Union
23

34
from ..models.record import Record
@@ -17,9 +18,7 @@ def set(
1718
transaction: Optional[Transaction] = None,
1819
) -> Dict[str, str]:
1920
"""Update a record by ID."""
20-
headers = Transaction._build_transaction_header(
21-
transaction.id if transaction else None
22-
)
21+
headers = Transaction._build_transaction_header(transaction)
2322
return self.client._make_request(
2423
"PUT", f"/api/v1/records/{record_id}", data, headers
2524
)
@@ -31,9 +30,8 @@ def update(
3130
transaction: Optional[Transaction] = None,
3231
) -> Dict[str, str]:
3332
"""Update a record by ID."""
34-
headers = Transaction._build_transaction_header(
35-
transaction.id if transaction else None
36-
)
33+
headers = Transaction._build_transaction_header(transaction)
34+
3735
return self.client._make_request(
3836
"PATCH", f"/api/v1/records/{record_id}", data, headers
3937
)
@@ -57,9 +55,7 @@ def create(
5755
Record object
5856
:param
5957
"""
60-
headers = Transaction._build_transaction_header(
61-
transaction.id if transaction else None
62-
)
58+
headers = Transaction._build_transaction_header(transaction)
6359

6460
payload = {
6561
"label": label,
@@ -89,9 +85,7 @@ def create_many(
8985
Returns:
9086
List of Record objects
9187
"""
92-
headers = Transaction._build_transaction_header(
93-
transaction.id if transaction else None
94-
)
88+
headers = Transaction._build_transaction_header(transaction)
9589

9690
payload = {
9791
"label": label,
@@ -106,39 +100,51 @@ def create_many(
106100
def attach(
107101
self,
108102
source: Union[str, Dict[str, Any]],
109-
target: Union[str, List[str], Dict[str, Any], List[Dict[str, Any]]],
103+
target: Union[
104+
str,
105+
List[str],
106+
Dict[str, Any],
107+
List[Dict[str, Any]],
108+
"Record",
109+
List["Record"],
110+
],
110111
options: Optional[RelationshipOptions] = None,
111112
transaction: Optional[Transaction] = None,
112113
) -> Dict[str, str]:
113114
"""Attach records to a source record."""
114-
headers = Transaction._build_transaction_header(
115-
transaction.id if transaction else None
116-
)
115+
headers = Transaction._build_transaction_header(transaction)
116+
117117
source_id = self._extract_target_ids(source)[0]
118118
target_ids = self._extract_target_ids(target)
119119
payload = {"targetIds": target_ids}
120120
if options:
121-
payload.update(options)
121+
payload.update(typing.cast(typing.Dict[str, typing.Any], options))
122122
return self.client._make_request(
123123
"POST", f"/api/v1/records/{source_id}/relations", payload, headers
124124
)
125125

126126
def detach(
127127
self,
128128
source: Union[str, Dict[str, Any]],
129-
target: Union[str, List[str], Dict[str, Any], List[Dict[str, Any]]],
129+
target: Union[
130+
str,
131+
List[str],
132+
Dict[str, Any],
133+
List[Dict[str, Any]],
134+
"Record",
135+
List["Record"],
136+
],
130137
options: Optional[RelationshipDetachOptions] = None,
131138
transaction: Optional[Transaction] = None,
132139
) -> Dict[str, str]:
133140
"""Detach records from a source record."""
134-
headers = Transaction._build_transaction_header(
135-
transaction.id if transaction else None
136-
)
141+
headers = Transaction._build_transaction_header(transaction)
142+
137143
source_id = self._extract_target_ids(source)[0]
138144
target_ids = self._extract_target_ids(target)
139145
payload = {"targetIds": target_ids}
140146
if options:
141-
payload.update(options)
147+
payload.update(typing.cast(typing.Dict[str, typing.Any], options))
142148
return self.client._make_request(
143149
"PUT", f"/api/v1/records/{source_id}/relations", payload, headers
144150
)
@@ -147,11 +153,13 @@ def delete(
147153
self, query: SearchQuery, transaction: Optional[Transaction] = None
148154
) -> Dict[str, str]:
149155
"""Delete records matching the query."""
150-
headers = Transaction._build_transaction_header(
151-
transaction.id if transaction else None
152-
)
156+
headers = Transaction._build_transaction_header(transaction)
157+
153158
return self.client._make_request(
154-
"PUT", "/api/v1/records/delete", query, headers
159+
"PUT",
160+
"/api/v1/records/delete",
161+
typing.cast(typing.Dict[str, typing.Any], query or {}),
162+
headers,
155163
)
156164

157165
def delete_by_id(
@@ -160,9 +168,8 @@ def delete_by_id(
160168
transaction: Optional[Transaction] = None,
161169
) -> Dict[str, str]:
162170
"""Delete records by ID(s)."""
163-
headers = Transaction._build_transaction_header(
164-
transaction.id if transaction else None
165-
)
171+
headers = Transaction._build_transaction_header(transaction)
172+
166173
if isinstance(id_or_ids, list):
167174
return self.client._make_request(
168175
"PUT",
@@ -183,16 +190,18 @@ def find(
183190
"""Find records matching the query."""
184191

185192
try:
186-
headers = Transaction._build_transaction_header(
187-
transaction.id if transaction else None
188-
)
193+
headers = Transaction._build_transaction_header(transaction)
194+
189195
path = (
190196
f"/api/v1/records/{record_id}/search"
191197
if record_id
192198
else "/api/v1/records/search"
193199
)
194200
response = self.client._make_request(
195-
"POST", path, data=query or {}, headers=headers
201+
"POST",
202+
path,
203+
data=typing.cast(typing.Dict[str, typing.Any], query or {}),
204+
headers=headers,
196205
)
197206
return [Record(self.client, record) for record in response.get("data")]
198207
except Exception:
@@ -206,9 +215,7 @@ def import_csv(
206215
transaction: Optional[Transaction] = None,
207216
) -> List[Dict[str, Any]]:
208217
"""Import data from CSV."""
209-
headers = Transaction._build_transaction_header(
210-
transaction.id if transaction else None
211-
)
218+
headers = Transaction._build_transaction_header(transaction)
212219

213220
payload = {
214221
"label": label,
@@ -222,15 +229,20 @@ def import_csv(
222229

223230
@staticmethod
224231
def _extract_target_ids(
225-
target: Union[str, List[str], Dict[str, Any], List[Dict[str, Any]]]
232+
target: Union[
233+
str,
234+
List[str],
235+
Dict[str, Any],
236+
List[Dict[str, Any]],
237+
"Record",
238+
List["Record"],
239+
]
226240
) -> List[str]:
227241
"""Extract target IDs from various input types."""
228242
if isinstance(target, str):
229243
return [target]
230244
elif isinstance(target, list):
231-
return [
232-
t["__id"] if isinstance(t, dict) and "__id" in t else t for t in target
233-
]
245+
return [t.get("__id", "") if isinstance(t, dict) else "" for t in target]
234246
elif isinstance(target, Record) and "__id" in target.data:
235247
return [target.data["__id"]]
236248
elif isinstance(target, dict) and "__id" in target:

‎src/rushdb/api/relationships.py

+6-4
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import typing
12
from typing import List, Optional, TypedDict, Union
23
from urllib.parse import urlencode
34

@@ -46,13 +47,14 @@ async def find(
4647
path = f"/records/relations/search{query_string}"
4748

4849
# Build headers with transaction if present
49-
headers = Transaction._build_transaction_header(
50-
transaction.id if transaction else None
51-
)
50+
headers = Transaction._build_transaction_header(transaction)
5251

5352
# Make request
5453
response = self.client._make_request(
55-
method="POST", path=path, data=query or {}, headers=headers
54+
method="POST",
55+
path=path,
56+
data=typing.cast(typing.Dict[str, typing.Any], query or {}),
57+
headers=headers,
5658
)
5759

5860
return response.data

‎src/rushdb/models/property.py

+3-16
Original file line numberDiff line numberDiff line change
@@ -18,24 +18,11 @@ class DatetimeObject(TypedDict, total=False):
1818

1919
DatetimeValue = Union[DatetimeObject, str]
2020
BooleanValue = bool
21-
NullValue = None
2221
NumberValue = float
2322
StringValue = str
2423

2524
# Property types
26-
PROPERTY_TYPE_BOOLEAN = "boolean"
27-
PROPERTY_TYPE_DATETIME = "datetime"
28-
PROPERTY_TYPE_NULL = "null"
29-
PROPERTY_TYPE_NUMBER = "number"
30-
PROPERTY_TYPE_STRING = "string"
31-
32-
PropertyType = Literal[
33-
PROPERTY_TYPE_BOOLEAN,
34-
PROPERTY_TYPE_DATETIME,
35-
PROPERTY_TYPE_NULL,
36-
PROPERTY_TYPE_NUMBER,
37-
PROPERTY_TYPE_STRING,
38-
]
25+
PropertyType = Literal["boolean", "datetime", "null", "number", "string"]
3926

4027

4128
class Property(TypedDict):
@@ -53,10 +40,10 @@ class PropertyWithValue(Property):
5340
value: Union[
5441
DatetimeValue,
5542
BooleanValue,
56-
NullValue,
43+
None,
5744
NumberValue,
5845
StringValue,
59-
List[Union[DatetimeValue, BooleanValue, NullValue, NumberValue, StringValue]],
46+
List[Union[DatetimeValue, BooleanValue, None, NumberValue, StringValue]],
6047
]
6148

6249

‎src/rushdb/models/record.py

+13-3
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,9 @@
1111
class Record:
1212
"""Represents a record in RushDB with methods for manipulation."""
1313

14-
def __init__(self, client: "RushDBClient", data: Dict[str, Any] = None):
14+
def __init__(
15+
self, client: "RushDBClient", data: Union[Dict[str, Any], None] = None
16+
):
1517
self._client = client
1618
# Handle different data formats
1719
if isinstance(data, dict):
@@ -25,7 +27,11 @@ def __init__(self, client: "RushDBClient", data: Dict[str, Any] = None):
2527
@property
2628
def id(self) -> str:
2729
"""Get record ID."""
28-
return self.data.get("__id")
30+
record_id = self.data.get("__id")
31+
if record_id is None:
32+
raise ValueError("Record ID is missing or None")
33+
34+
return record_id
2935

3036
@property
3137
def proptypes(self) -> str:
@@ -40,7 +46,11 @@ def label(self) -> str:
4046
@property
4147
def timestamp(self) -> int:
4248
"""Get record timestamp from ID."""
43-
parts = self.data.get("__id").split("-")
49+
record_id = self.data.get("__id")
50+
if record_id is None:
51+
raise ValueError("Record ID is missing or None")
52+
53+
parts = record_id.split("-")
4454
high_bits_hex = parts[0] + parts[1][:4]
4555
return int(high_bits_hex, 16)
4656

0 commit comments

Comments
 (0)