Skip to content

Commit 72965ce

Browse files
authored
Merge pull request #31 from jymchng/fix-enum-assignment
Fix enum assignment
2 parents 39a9a10 + 3bb3419 commit 72965ce

File tree

10 files changed

+338
-39
lines changed

10 files changed

+338
-39
lines changed

Makefile

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -124,9 +124,12 @@ test-str:
124124
test-async:
125125
$(PYTHON) -m pytest tests/test_async.py $(PYTEST_FLAGS)
126126

127+
test-enums:
128+
$(PYTHON) -m pytest tests/test_newtype_enums.py $(PYTEST_FLAGS)
129+
127130
# Run memory leak tests
128131
test-leak:
129-
$(PYTHON) -m pytest --enable-leak-tracking -W error --stacks 10 tests/test_newtype_init.py $(PYTEST_FLAGS)
132+
$(PYTHON) -m pytest --enable-leak-tracking -W error tests/test_newtype_init.py $(PYTEST_FLAGS)
130133

131134
# Run a specific test file (usage: make test-file FILE=test_newtype.py)
132135
test-file:

examples/newtype_enums.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ class ENV(NewType(str), Enum): # type: ignore[misc]
1414
PREPROD = "PREPROD"
1515
PROD = "PROD"
1616

17+
1718
class RegularENV(str, Enum):
1819

1920
LOCAL = "LOCAL"
@@ -58,6 +59,7 @@ class RollYourOwnNewTypeEnum(ENVVariant, Enum): # type: ignore[no-redef]
5859
PREPROD = "PREPROD"
5960
PROD = "PROD"
6061

62+
6163
# mypy doesn't raise errors here
6264
def test_nt_env_replace() -> None:
6365

examples/newtype_enums_int.py

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
from enum import Enum
2+
from weakref import WeakValueDictionary
3+
4+
import pytest
5+
6+
from newtype import NewType
7+
8+
9+
class GenericWrappedBoundedInt(NewType(int)):
10+
MAX_VALUE: int = 0
11+
12+
__CONCRETE_BOUNDED_INTS__ = WeakValueDictionary()
13+
14+
def __new__(cls, value: int):
15+
inst = super().__new__(cls, value % cls.MAX_VALUE)
16+
return inst
17+
18+
def __repr__(self) -> str:
19+
return f"<BoundedInt[MAX_VALUE={self.MAX_VALUE}]: {super().__repr__()}>"
20+
21+
def __str__(self) -> str:
22+
return str(int(self))
23+
24+
def __class_getitem__(cls, idx=MAX_VALUE):
25+
if not isinstance(idx, int):
26+
raise TypeError(f"cannot make `BoundedInt[{idx}]`")
27+
28+
if idx not in cls.__CONCRETE_BOUNDED_INTS__:
29+
30+
class ConcreteBoundedInt(cls):
31+
MAX_VALUE = idx
32+
33+
cls.__CONCRETE_BOUNDED_INTS__[idx] = ConcreteBoundedInt
34+
35+
return cls.__CONCRETE_BOUNDED_INTS__[idx]
36+
37+
38+
class Severity(GenericWrappedBoundedInt[5], Enum):
39+
DEBUG = 0
40+
INFO = 1
41+
WARNING = 2
42+
ERROR = 3
43+
CRITICAL = 4
44+
45+
46+
def test_severity():
47+
assert Severity.DEBUG == 0
48+
assert Severity.INFO == 1
49+
assert Severity.WARNING == 2
50+
assert Severity.ERROR == 3
51+
assert Severity.CRITICAL == 4
52+
53+
with pytest.raises(AttributeError, match=r"[c|C]annot\s+reassign\s+\w+"):
54+
Severity.ERROR += 1
55+
56+
severity = Severity.ERROR
57+
assert severity == 3
58+
59+
severity += 1
60+
assert severity == 4
61+
assert severity != 3
62+
assert isinstance(severity, int)
63+
assert isinstance(severity, Severity)
64+
assert severity is not Severity.ERROR
65+
assert severity is Severity.CRITICAL
66+
67+
severity -= 1
68+
assert severity == 3
69+
assert severity != 4
70+
assert isinstance(severity, int)
71+
assert isinstance(severity, Severity)
72+
assert severity is Severity.ERROR
73+
assert severity is not Severity.CRITICAL
74+
75+
severity = Severity.DEBUG
76+
assert severity == 0
77+
assert str(severity.value) == "0"
78+
with pytest.raises(ValueError, match=r"\d+ is not a valid Severity"):
79+
severity -= 1
80+
81+
severity = Severity.CRITICAL
82+
assert severity == 4
83+
assert str(severity.value) == "4"
84+
with pytest.raises(ValueError, match=r"\d+ is not a valid Severity"):
85+
severity += 1

newtype/extensions/newtype_init.c

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,7 @@ static PyObject* NewTypeInit_call(NewTypeInitObject* self,
101101
PyObject* func;
102102

103103
if (self->has_get) {
104+
DEBUG_PRINT("`self->has_get`: %d\n", self->has_get);
104105
if (self->obj == NULL && self->cls == NULL) {
105106
// free standing function
106107
PyErr_SetString(
@@ -117,6 +118,8 @@ static PyObject* NewTypeInit_call(NewTypeInitObject* self,
117118
self->func_get, self->obj, self->cls, NULL);
118119
}
119120
} else {
121+
DEBUG_PRINT("`self->func_get`: %s\n",
122+
PyUnicode_AsUTF8(PyObject_Repr(self->func_get)));
120123
func = self->func_get;
121124
}
122125

@@ -179,6 +182,7 @@ static PyObject* NewTypeInit_call(NewTypeInitObject* self,
179182
result = PyObject_Call(func, args, kwds);
180183
} else {
181184
PyErr_SetString(PyExc_TypeError, "Invalid type object in descriptor");
185+
DEBUG_PRINT("`self->cls` is not a valid type object\n");
182186
result = NULL;
183187
}
184188

newtype/extensions/newtype_meth.c

Lines changed: 102 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,8 @@ static PyObject* NewTypeMethod_call(NewTypeMethodObject* self,
9494
self->func_get, Py_None, self->wrapped_cls, NULL);
9595
} else {
9696
DEBUG_PRINT("`self->obj` is not NULL\n");
97+
DEBUG_PRINT("`self->wrapped_cls`: %s\n",
98+
PyUnicode_AsUTF8(PyObject_Repr(self->wrapped_cls)));
9799
func = PyObject_CallFunctionObjArgs(
98100
self->func_get, self->obj, self->wrapped_cls, NULL);
99101
}
@@ -145,7 +147,9 @@ static PyObject* NewTypeMethod_call(NewTypeMethodObject* self,
145147
{ // now we try to build an instance of the subtype
146148
DEBUG_PRINT("`result` is an instance of `self->wrapped_cls`\n");
147149
PyObject *init_args, *init_kwargs;
148-
PyObject *new_inst, *args_combined;
150+
PyObject* new_inst;
151+
PyObject* args_combined = NULL;
152+
Py_ssize_t args_len = 0;
149153

150154
if (self->obj == NULL) {
151155
PyObject* first_elem;
@@ -158,73 +162,136 @@ static PyObject* NewTypeMethod_call(NewTypeMethodObject* self,
158162
first_elem = PyTuple_GetItem(args, 0);
159163
Py_XINCREF(
160164
first_elem); // Increment reference count of the first element
161-
165+
DEBUG_PRINT("`first_elem`: %s\n",
166+
PyUnicode_AsUTF8(PyObject_Repr(first_elem)));
162167
} else { // `args` is empty here, then we are done actually
168+
DEBUG_PRINT("`args` is empty\n");
163169
goto done;
164170
};
165171
if (PyObject_IsInstance(first_elem, (PyObject*)self->cls)) {
166172
init_args = PyObject_GetAttrString(first_elem, NEWTYPE_INIT_ARGS_STR);
167173
init_kwargs =
168174
PyObject_GetAttrString(first_elem, NEWTYPE_INIT_KWARGS_STR);
175+
DEBUG_PRINT("`init_args`: %s\n",
176+
PyUnicode_AsUTF8(PyObject_Repr(init_args)));
177+
DEBUG_PRINT("`init_kwargs`: %s\n",
178+
PyUnicode_AsUTF8(PyObject_Repr(init_kwargs)));
169179
} else { // first element is not the subtype, so we are done also
180+
DEBUG_PRINT("`first_elem` is not the subtype\n");
170181
goto done;
171182
}
172183
Py_XDECREF(first_elem);
173184
} else { // `self->obj` is not NULL
174185

186+
DEBUG_PRINT("`self->obj` is not NULL\n");
175187
init_args = PyObject_GetAttrString(self->obj, NEWTYPE_INIT_ARGS_STR);
176188
init_kwargs = PyObject_GetAttrString(self->obj, NEWTYPE_INIT_KWARGS_STR);
189+
DEBUG_PRINT("`init_args`: %s\n",
190+
PyUnicode_AsUTF8(PyObject_Repr(init_args)));
191+
DEBUG_PRINT("`init_kwargs`: %s\n",
192+
PyUnicode_AsUTF8(PyObject_Repr(init_kwargs)));
177193
}
178194

179-
Py_ssize_t args_len = PyTuple_Size(init_args);
180-
Py_ssize_t combined_args_len = 1 + args_len;
181-
args_combined = PyTuple_New(combined_args_len);
182-
if (args_combined == NULL) {
183-
Py_XDECREF(init_args);
184-
Py_XDECREF(init_kwargs);
185-
Py_DECREF(result);
186-
return NULL; // Use return NULL instead of Py_RETURN_NONE
187-
}
188-
189-
// Set the first item of the new tuple to `result`
190-
PyTuple_SET_ITEM(args_combined,
191-
0,
192-
result); // `result` is now owned by `args_combined`
193-
194-
// Copy items from `init_args` to `args_combined`
195-
for (Py_ssize_t i = 0; i < args_len; i++) {
196-
PyObject* item = PyTuple_GetItem(init_args, i); // Borrowed reference
197-
if (item == NULL) {
198-
Py_DECREF(args_combined);
195+
if (init_args != NULL) {
196+
DEBUG_PRINT("`init_args` is not NULL\n");
197+
args_len = PyTuple_Size(init_args);
198+
DEBUG_PRINT("`args_len`: %zd\n", args_len);
199+
Py_ssize_t combined_args_len = 1 + args_len;
200+
DEBUG_PRINT("`combined_args_len`: %zd\n", combined_args_len);
201+
args_combined = PyTuple_New(combined_args_len);
202+
DEBUG_PRINT("`args_combined`: %s\n",
203+
PyUnicode_AsUTF8(PyObject_Repr(args_combined)));
204+
if (args_combined == NULL) {
199205
Py_XDECREF(init_args);
200206
Py_XDECREF(init_kwargs);
201-
return NULL;
207+
Py_DECREF(result);
208+
DEBUG_PRINT("`args_combined` is NULL\n");
209+
return NULL; // Use return NULL instead of Py_RETURN_NONE
202210
}
203-
Py_INCREF(item); // Increase reference count
211+
// Set the first item of the new tuple to `result`
204212
PyTuple_SET_ITEM(args_combined,
205-
i + 1,
206-
item); // `item` is now owned by `args_combined`
213+
0,
214+
result); // `result` is now owned by `args_combined`
215+
216+
// Copy items from `init_args` to `args_combined`
217+
for (Py_ssize_t i = 0; i < args_len; i++) {
218+
PyObject* item = PyTuple_GetItem(init_args, i); // Borrowed reference
219+
if (item == NULL) {
220+
DEBUG_PRINT("`item` is NULL\n");
221+
Py_DECREF(args_combined);
222+
Py_XDECREF(init_args);
223+
Py_XDECREF(init_kwargs);
224+
return NULL;
225+
}
226+
DEBUG_PRINT("`item`: %s\n", PyUnicode_AsUTF8(PyObject_Repr(item)));
227+
Py_INCREF(item); // Increase reference count
228+
PyTuple_SET_ITEM(args_combined,
229+
i + 1,
230+
item); // `item` is now owned by `args_combined`
231+
}
232+
DEBUG_PRINT("`args_combined`: %s\n",
233+
PyUnicode_AsUTF8(PyObject_Repr(args_combined)));
207234
}
208-
DEBUG_PRINT("`args_combined`: %s\n",
209-
PyUnicode_AsUTF8(PyObject_Repr(args_combined)));
235+
236+
if (init_args == NULL || init_kwargs == NULL) {
237+
DEBUG_PRINT("`init_args` or `init_kwargs` is NULL\n");
238+
};
210239

211240
if (init_kwargs != NULL) {
212241
DEBUG_PRINT("`init_kwargs`: %s\n",
213242
PyUnicode_AsUTF8(PyObject_Repr(init_kwargs)));
214243
};
215244

216-
// Call the function or constructor
245+
// If `args_combined` is NULL, create a new tuple with one item
246+
// and set `result` as the first item of the tuple
247+
if (init_args == NULL) {
248+
DEBUG_PRINT("`init_args` is NULL\n");
249+
250+
if (PyObject_SetAttrString(
251+
self->obj, NEWTYPE_INIT_ARGS_STR, PyTuple_New(0))
252+
< 0)
253+
{
254+
result = NULL;
255+
goto done;
256+
}
257+
if (PyObject_SetAttrString(
258+
self->obj, NEWTYPE_INIT_KWARGS_STR, PyDict_New())
259+
< 0)
260+
{
261+
result = NULL;
262+
goto done;
263+
}
264+
265+
args_combined = PyTuple_New(1); // Allocate tuple with one element
266+
Py_INCREF(result);
267+
PyTuple_SET_ITEM(args_combined, 0, result);
268+
DEBUG_PRINT("`args_combined`: %s\n",
269+
PyUnicode_AsUTF8(PyObject_Repr(args_combined)));
270+
new_inst =
271+
PyObject_Call((PyObject*)self->cls, args_combined, init_kwargs);
272+
if (new_inst == NULL) {
273+
DEBUG_PRINT("`new_inst` is NULL\n");
274+
Py_DECREF(result);
275+
Py_DECREF(self->obj);
276+
Py_DECREF(args_combined);
277+
return NULL;
278+
}
279+
Py_DECREF(result);
280+
Py_DECREF(self->obj);
281+
Py_DECREF(args_combined);
282+
DEBUG_PRINT("`new_inst`: %s\n",
283+
PyUnicode_AsUTF8(PyObject_Repr(new_inst)));
284+
return new_inst;
285+
}
286+
217287
new_inst = PyObject_Call((PyObject*)self->cls, args_combined, init_kwargs);
218288

219289
// Clean up
220-
Py_DECREF(args_combined); // Decrement reference count of `args_combined`
290+
Py_XDECREF(args_combined); // Decrement reference count of `args_combined`
221291
Py_XDECREF(init_args);
222292
Py_XDECREF(init_kwargs);
223293

224-
// Ensure proper error propagation
225-
if (new_inst == NULL) {
226-
return NULL;
227-
}
294+
DEBUG_PRINT("`new_inst`: %s\n", PyUnicode_AsUTF8(PyObject_Repr(new_inst)));
228295

229296
// Only proceed if we have all required objects and dictionaries
230297
if (self->obj != NULL && result != NULL && new_inst != NULL
@@ -427,6 +494,7 @@ static PyObject* NewTypeMethod_call(NewTypeMethodObject* self,
427494

428495
done:
429496
Py_XINCREF(result);
497+
DEBUG_PRINT("DONE! `result`: %s\n", PyUnicode_AsUTF8(PyObject_Repr(result)));
430498
return result;
431499
}
432500

newtype/newtype.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -180,7 +180,7 @@ class BaseNewType(base_type): # type: ignore[valid-type, misc]
180180

181181
if hasattr(base_type, "__slots__"):
182182
__slots__ = (
183-
*base_type.__slots__,
183+
# *base_type.__slots__,
184184
NEWTYPE_INIT_ARGS_STR,
185185
NEWTYPE_INIT_KWARGS_STR,
186186
)
@@ -224,10 +224,14 @@ def __init_subclass__(cls, **init_subclass_context: Any) -> None:
224224
and not func_is_excluded(v)
225225
):
226226
setattr(cls, k, NewTypeMethod(v, base_type))
227+
227228
else:
228229
if k == "__dict__":
229230
continue
230-
setattr(cls, k, v)
231+
try:
232+
setattr(cls, k, v)
233+
except AttributeError:
234+
continue
231235
cls.__init__ = NewTypeInit(constructor) # type: ignore[method-assign]
232236

233237
def __new__(cls, value: Any = None, *_args: Any, **_kwargs: Any) -> "BaseNewType":

pyproject.toml

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ build-backend = "poetry_dynamic_versioning.backend"
44

55
[tool.poetry]
66
name = "python-newtype"
7-
version = "0.1.5"
7+
version = "0.1.6"
88
homepage = "https://github.com/jymchng/python-newtype-dev"
99
repository = "https://github.com/jymchng/python-newtype-dev"
1010
license = "MIT"
@@ -146,6 +146,7 @@ exclude = [
146146
"examples/newtype_enums.py",
147147
"examples/mutable.py",
148148
"examples/pydantic-compat.py",
149+
"examples/newtype_enums_int.py",
149150
]
150151

151152
[tool.ruff.format]

tests/build_test_pyvers_docker_images.sh

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22

33
# Create logs directory if it doesn't exist
44
mkdir -p ./tests/logs
5-
make build
65

76
# Build Docker images in parallel with logging
87
docker build -t python-newtype-test-mul-vers:3.8 -f ./tests/Dockerfile-test-py3.8 . > ./tests/logs/py3.8-test.log 2>&1 &

0 commit comments

Comments
 (0)