Skip to content

Commit c21e8be

Browse files
committed
Housekeeping
- Refactored the main module - Increased code coverage
1 parent e26f979 commit c21e8be

6 files changed

Lines changed: 513 additions & 88 deletions

File tree

apywire/compiler.py

Lines changed: 103 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -128,66 +128,8 @@ def _compile_property(
128128
# lambda passed to a thread pool executor is pure synchronous.
129129
pre_statements: list[ast.stmt] = []
130130

131-
# Helper to extract `await self.<name>()` occurrences from AST
132-
# expressions and replace them with local variables. This lets us
133-
# run a synchronous lambda in `run_in_executor` without `await`
134-
# nodes embedded inside it.
135-
counter = 0
136-
137-
def _replace_awaits_with_locals(node: ast.expr) -> ast.expr:
138-
nonlocal counter
139-
140-
if isinstance(node, ast.Await):
141-
inner = node.value
142-
if (
143-
isinstance(inner, ast.Call)
144-
and isinstance(inner.func, ast.Attribute)
145-
and isinstance(inner.func.value, ast.Name)
146-
and inner.func.value.id == "self"
147-
):
148-
# produce a unique variable name and precompute the await
149-
counter += 1
150-
var_name = f"{COMPILED_VAR_PREFIX}{counter}"
151-
assign = ast.Assign(
152-
targets=[ast.Name(id=var_name, ctx=ast.Store())],
153-
value=node,
154-
)
155-
pre_statements.append(assign)
156-
return ast.Name(id=var_name, ctx=ast.Load())
157-
return node
158-
159-
# Recurse into common composite nodes
160-
if isinstance(node, ast.Call):
161-
new_args = [_replace_awaits_with_locals(a) for a in node.args]
162-
new_keywords = [
163-
ast.keyword(
164-
arg=k.arg, value=_replace_awaits_with_locals(k.value)
165-
)
166-
for k in node.keywords
167-
]
168-
return ast.Call(
169-
func=node.func, args=new_args, keywords=new_keywords
170-
)
171-
if isinstance(node, ast.Dict):
172-
new_keys = [
173-
_replace_awaits_with_locals(k) if k is not None else None
174-
for k in node.keys
175-
]
176-
new_values = [
177-
_replace_awaits_with_locals(v) for v in node.values
178-
]
179-
return ast.Dict(keys=new_keys, values=new_values)
180-
if isinstance(node, ast.List):
181-
return ast.List(
182-
elts=[_replace_awaits_with_locals(e) for e in node.elts],
183-
ctx=node.ctx,
184-
)
185-
if isinstance(node, ast.Tuple):
186-
return ast.Tuple(
187-
elts=[_replace_awaits_with_locals(e) for e in node.elts],
188-
ctx=node.ctx,
189-
)
190-
return node
131+
# Counter for generating unique variable names (mutable list)
132+
counter = [0]
191133

192134
args: list[ast.expr] = []
193135

@@ -210,7 +152,9 @@ def _replace_awaits_with_locals(node: ast.expr) -> ast.expr:
210152
for i, value in enumerate(args_data):
211153
raw_val_ast = self._astify(value, aio=aio)
212154
if aio:
213-
val_ast = _replace_awaits_with_locals(raw_val_ast)
155+
val_ast = self._replace_awaits_with_locals(
156+
raw_val_ast, pre_statements, counter
157+
)
214158
var_name = f"{COMPILED_ARG_PREFIX}{i}"
215159
assign = ast.Assign(
216160
targets=[ast.Name(id=var_name, ctx=ast.Store())],
@@ -229,7 +173,9 @@ def _replace_awaits_with_locals(node: ast.expr) -> ast.expr:
229173
# Replace any awaited accessors with local precomputed
230174
# variables and assign all values to local variables so
231175
# that the executor lambda can be synchronous.
232-
val_ast = _replace_awaits_with_locals(raw_val_ast)
176+
val_ast = self._replace_awaits_with_locals(
177+
raw_val_ast, pre_statements, counter
178+
)
233179
# Use named variables for top-level keyword args to make
234180
# the generated code more readable and deterministic.
235181
var_name = f"{COMPILED_VAR_PREFIX}{key}"
@@ -494,6 +440,101 @@ def _replace_awaits_with_locals(node: ast.expr) -> ast.expr:
494440
)
495441
return func_def
496442

443+
def _replace_awaits_with_locals(
444+
self,
445+
node: ast.expr,
446+
pre_statements: list[ast.stmt],
447+
counter: list[int],
448+
) -> ast.expr:
449+
"""Replace await expressions with local variable assignments.
450+
451+
Recursively processes AST nodes to replace `ast.Await` expressions
452+
with local variable assignments. This ensures that synchronous lambdas
453+
executed in a thread pool do not contain `await` nodes.
454+
455+
Args:
456+
node: AST expression node to process
457+
pre_statements: List to accumulate pre-computation statements
458+
counter: Single-element list used as mutable counter for var naming
459+
460+
Returns:
461+
Transformed AST expression node
462+
"""
463+
if isinstance(node, ast.Await):
464+
inner = node.value
465+
if (
466+
isinstance(inner, ast.Call)
467+
and isinstance(inner.func, ast.Attribute)
468+
and isinstance(inner.func.value, ast.Name)
469+
and inner.func.value.id == "self"
470+
):
471+
# produce a unique variable name and precompute the await
472+
counter[0] += 1
473+
var_name = f"{COMPILED_VAR_PREFIX}{counter[0]}"
474+
assign = ast.Assign(
475+
targets=[ast.Name(id=var_name, ctx=ast.Store())],
476+
value=node,
477+
)
478+
pre_statements.append(assign)
479+
return ast.Name(id=var_name, ctx=ast.Load())
480+
return node
481+
482+
# Recurse into common composite nodes
483+
if isinstance(node, ast.Call):
484+
new_args = [
485+
self._replace_awaits_with_locals(a, pre_statements, counter)
486+
for a in node.args
487+
]
488+
new_keywords = [
489+
ast.keyword(
490+
arg=k.arg,
491+
value=self._replace_awaits_with_locals(
492+
k.value, pre_statements, counter
493+
),
494+
)
495+
for k in node.keywords
496+
]
497+
return ast.Call(
498+
func=node.func, args=new_args, keywords=new_keywords
499+
)
500+
if isinstance(node, ast.Dict):
501+
new_keys = [
502+
(
503+
self._replace_awaits_with_locals(
504+
k, pre_statements, counter
505+
)
506+
if k is not None
507+
else None
508+
)
509+
for k in node.keys
510+
]
511+
new_values = [
512+
self._replace_awaits_with_locals(v, pre_statements, counter)
513+
for v in node.values
514+
]
515+
return ast.Dict(keys=new_keys, values=new_values)
516+
if isinstance(node, ast.List):
517+
return ast.List(
518+
elts=[
519+
self._replace_awaits_with_locals(
520+
e, pre_statements, counter
521+
)
522+
for e in node.elts
523+
],
524+
ctx=node.ctx,
525+
)
526+
if isinstance(node, ast.Tuple):
527+
return ast.Tuple(
528+
elts=[
529+
self._replace_awaits_with_locals(
530+
e, pre_statements, counter
531+
)
532+
for e in node.elts
533+
],
534+
ctx=node.ctx,
535+
)
536+
return node
537+
497538
def _compile_constant_property(
498539
self,
499540
name: str,

apywire/threads.py

Lines changed: 49 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,38 @@ def _release_held_locks(self) -> None:
8888
return
8989
for lock in reversed(self._local.held_locks):
9090
lock.release()
91-
self._local.held_locks = []
91+
self._local.held_locks.clear()
92+
93+
def _check_cache(self, name: str) -> tuple[bool, object | None]:
94+
"""Check if attribute is cached, return (found, value) tuple.
95+
96+
Checks both _values dict (if present) and _<name> attribute.
97+
Returns (True, value) if cached, (False, None) if not found.
98+
"""
99+
cache_attr = f"{CACHE_ATTR_PREFIX}{name}"
100+
101+
# Check _values dict first (runtime path)
102+
if hasattr(self, "_values"):
103+
values_dict = cast(dict[str, object], getattr(self, "_values"))
104+
if name in values_dict:
105+
return (True, values_dict[name])
106+
107+
# Check cached attribute (compiled path)
108+
if hasattr(self, cache_attr):
109+
return (True, cast(object, getattr(self, cache_attr)))
110+
111+
return (False, None)
112+
113+
def _set_cache(self, name: str, value: object) -> None:
114+
"""Store value in cache using appropriate storage mechanism.
115+
116+
Uses _values dict if present, otherwise uses _<name> attribute.
117+
"""
118+
if hasattr(self, "_values"):
119+
self._values[name] = value
120+
else:
121+
cache_attr = f"{CACHE_ATTR_PREFIX}{name}"
122+
setattr(self, cache_attr, value)
92123

93124
def _wrap_instantiation_error(self, e: Exception, name: str) -> NoReturn:
94125
"""Wrap exceptions during instantiation with proper context.
@@ -119,12 +150,10 @@ def _instantiate_attr(
119150
thread-safe locking logic permits it to run. Returns the cached
120151
attribute (``self._<name>``) if already present.
121152
"""
122-
cache_attr = f"{CACHE_ATTR_PREFIX}{name}"
123-
# check cache mapping or attribute
124-
if hasattr(self, "_values") and name in self._values:
125-
return self._values[name]
126-
if hasattr(self, cache_attr):
127-
return cast(object, getattr(self, cache_attr))
153+
# Check cache first
154+
found, value = self._check_cache(name)
155+
if found:
156+
return value
128157

129158
lock = self._get_attribute_lock(name)
130159
mode: str | None = getattr(self._local, "mode", None)
@@ -134,8 +163,9 @@ def _instantiate_attr(
134163
# If not, fall back to global lock.
135164
if lock.acquire(blocking=False):
136165
try:
137-
if hasattr(self, cache_attr):
138-
return cast(object, getattr(self, cache_attr))
166+
found, value = self._check_cache(name)
167+
if found:
168+
return value
139169
# Enter optimistic mode and track held locks to allow
140170
# nested instantiation to detect per-attribute
141171
# conflicts.
@@ -155,11 +185,8 @@ def _instantiate_attr(
155185
lock.release()
156186
self._wrap_instantiation_error(e, name)
157187
else:
158-
# store both in mapping and attribute when available
159-
if hasattr(self, "_values"):
160-
self._values[name] = inst
161-
else:
162-
setattr(self, cache_attr, inst)
188+
# Store in cache
189+
self._set_cache(name, inst)
163190
# Release the optimistic held locks before returning
164191
self._release_held_locks()
165192
return inst
@@ -185,13 +212,11 @@ def _instantiate_attr(
185212
attempts = 0
186213
while True:
187214
try:
188-
if hasattr(self, cache_attr):
189-
return cast(object, getattr(self, cache_attr))
215+
found, value = self._check_cache(name)
216+
if found:
217+
return value
190218
inst = maker()
191-
if hasattr(self, "_values"):
192-
self._values[name] = inst
193-
else:
194-
setattr(self, cache_attr, inst)
219+
self._set_cache(name, inst)
195220
return inst
196221
except LockUnavailableError:
197222
attempts += 1
@@ -219,13 +244,11 @@ def _instantiate_attr(
219244
held = self._get_held_locks()
220245
held.append(lock)
221246
try:
222-
if hasattr(self, cache_attr):
223-
return cast(object, getattr(self, cache_attr))
247+
found, value = self._check_cache(name)
248+
if found:
249+
return value
224250
inst = maker()
225-
if hasattr(self, "_values"):
226-
self._values[name] = inst
227-
else:
228-
setattr(self, cache_attr, inst)
251+
self._set_cache(name, inst)
229252
return inst
230253
finally:
231254
# Locks are intentionally not released here; they will be

tests/test_compile_aio.py

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -284,3 +284,83 @@ async def run() -> None:
284284
assert val == complex(1.0, 2.0)
285285

286286
asyncio.run(run())
287+
288+
289+
def test_aio_compile_await_non_self_call() -> None:
290+
"""Test _replace_awaits_with_locals with non-self await nodes."""
291+
from apywire import Spec, WiringCompiler
292+
293+
# Create spec with async compilation that could have nested awaits
294+
spec: Spec = {
295+
"datetime.datetime now": {"year": 2025, "month": 1, "day": 1},
296+
"datetime.datetime later": {"year": 2025, "month": 6, "day": 15},
297+
}
298+
299+
code = WiringCompiler(spec, thread_safe=False).compile(aio=True)
300+
301+
# Verify code was generated (exact await handling is complex)
302+
assert "async def now(self):" in code
303+
assert "async def later(self):" in code
304+
305+
306+
def test_aio_compile_nested_call_replacement() -> None:
307+
"""Test _replace_awaits_with_locals in Call nodes."""
308+
import sys
309+
from types import ModuleType
310+
311+
from apywire import Spec, WiringCompiler
312+
313+
class MockMod(ModuleType):
314+
def combine(
315+
self, *args: object, **kwargs: object
316+
) -> dict[str, object]:
317+
return {"args": args, **kwargs}
318+
319+
sys.modules["nested_call"] = MockMod("nested_call")
320+
321+
try:
322+
spec: Spec = {
323+
"nested_call.combine root": {
324+
"x": "{a}",
325+
"y": "{b}",
326+
},
327+
"nested_call.combine a": {},
328+
"nested_call.combine b": {},
329+
}
330+
code = WiringCompiler(spec, thread_safe=False).compile(aio=True)
331+
332+
# Verify async code was generated with replaced awaits
333+
assert "async def root(self):" in code
334+
assert "__val_" in code
335+
336+
finally:
337+
del sys.modules["nested_call"]
338+
339+
340+
def test_compile_constant_accessor_skip() -> None:
341+
"""Test that constant accessors skip already-parsed items."""
342+
import sys
343+
from types import ModuleType
344+
345+
from apywire import Spec, WiringCompiler
346+
347+
class MockMod(ModuleType):
348+
pass
349+
350+
sys.modules["const_test"] = MockMod("const_test")
351+
352+
try:
353+
# Mix of parsed wired objects and constants
354+
spec: Spec = {
355+
"datetime.datetime obj": {"year": 2025, "month": 1, "day": 1},
356+
"const": "value",
357+
}
358+
code = WiringCompiler(spec, thread_safe=False).compile()
359+
360+
# Both should be in code
361+
assert "def obj(self):" in code
362+
assert "def const(self):" in code
363+
assert "return 'value'" in code
364+
365+
finally:
366+
del sys.modules["const_test"]

0 commit comments

Comments
 (0)