jmux 0.0.5__py3-none-any.whl → 0.1.0__py3-none-any.whl

jmux/__init__.py

@@ -0,0 +1,11 @@
+from jmux.awaitable import AwaitableValue, StreamableValues
+from jmux.base import StreamableBaseModel, Streamed
+from jmux.demux import JMux
+
+__all__ = [
+    "AwaitableValue",
+    "JMux",
+    "Streamed",
+    "StreamableBaseModel",
+    "StreamableValues",
+]

jmux/awaitable.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from asyncio import Event, Queue
+import math
 from enum import Enum
 from types import NoneType
 from typing import (
@@ -14,6 +14,9 @@ from typing import (
     runtime_checkable,
 )
 
+from anyio import Event, create_memory_object_stream
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+
 from jmux.error import NothingEmittedError, SinkClosedError
 from jmux.helpers import extract_types_from_generic_alias
 
@@ -27,7 +30,7 @@ class SinkType(Enum):
 
 class UnderlyingGenericMixin(Generic[T]):
     """
-    A mixin class that provides methods for inspecting the generic types of a 
+    A mixin class that provides methods for inspecting the generic types of a
     class at runtime.
     """
 
@@ -105,12 +108,17 @@ class IAsyncSink(Protocol[T]):
 class StreamableValues(UnderlyingGenericMixin[T], Generic[T]):
     """
     A class that represents a stream of values that can be asynchronously iterated over.
-    It uses an asyncio.Queue to store the items and allows for putting items into the
-    stream and closing it when no more items will be added.
+    It uses anyio memory object streams to store the items and allows for putting items
+    into the stream and closing it when no more items will be added.
     """
 
+    _send_stream: MemoryObjectSendStream[T | None]
+    _receive_stream: MemoryObjectReceiveStream[T | None]
+
     def __init__(self):
-        self._queue = Queue[T | None]()
+        self._send_stream, self._receive_stream = create_memory_object_stream[T | None](
+            max_buffer_size=math.inf
+        )
         self._last_item: T | None = None
         self._closed = False
 
@@ -142,7 +150,7 @@ class StreamableValues(UnderlyingGenericMixin[T], Generic[T]):
         if self._closed:
             raise ValueError("Cannot put item into a closed sink.")
         self._last_item = item
-        await self._queue.put(item)
+        await self._send_stream.send(item)
 
     async def close(self):
         """
@@ -154,10 +162,11 @@ class StreamableValues(UnderlyingGenericMixin[T], Generic[T]):
         if self._closed:
             raise SinkClosedError(
                 f"SinkType {self.get_sink_type()}[{self.get_underlying_main_generic()}]"
-                  + " is already closed."
+                + " is already closed."
             )
         self._closed = True
-        await self._queue.put(None)
+        await self._send_stream.send(None)
+        await self._send_stream.aclose()
 
     async def ensure_closed(self):
         """
@@ -195,13 +204,16 @@ class StreamableValues(UnderlyingGenericMixin[T], Generic[T]):
         return self._stream()
 
     async def _stream(self) -> AsyncGenerator[T, None]:
-        while True:
-            item = await self._queue.get()
-            if item is None and self._closed:
-                break
-            if item is None:
-                raise ValueError("Received None item, but the sink is not closed.")
-            yield item
+        try:
+            while True:
+                item = await self._receive_stream.receive()
+                if item is None and self._closed:
+                    break
+                if item is None:
+                    raise ValueError("Received None item, but the sink is not closed.")
+                yield item
+        finally:
+            await self._receive_stream.aclose()
 
 
 class AwaitableValue(UnderlyingGenericMixin[T], Generic[T]):
@@ -242,7 +254,7 @@ class AwaitableValue(UnderlyingGenericMixin[T], Generic[T]):
         if self._is_closed:
             raise SinkClosedError(
                 f"SinkType {self.get_sink_type()}"
-                +"[{self.get_underlying_main_generic().__name__}] is already closed."
+                + "[{self.get_underlying_main_generic().__name__}] is already closed."
             )
         elif not self._event.is_set() and NoneType in self.get_underlying_generics():
             self._event.set()

jmux/base.py

@@ -0,0 +1,9 @@
+from pydantic import BaseModel
+
+
+class StreamableBaseModel(BaseModel):
+    pass
+
+
+class Streamed:
+    pass

jmux/cli.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+from jmux.generator import find_streamable_models, generate_jmux_code
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        prog="jmux",
+        description="JMux CLI for generating JMux classes from Pydantic models",
+    )
+    subparsers = parser.add_subparsers(dest="command")
+
+    gen_parser = subparsers.add_parser(
+        "generate",
+        help="Generate JMux classes from StreamableBaseModel subclasses",
+    )
+    gen_parser.add_argument(
+        "--root",
+        type=Path,
+        default=Path("."),
+        help="Root directory to scan for StreamableBaseModel subclasses",
+    )
+
+    args = parser.parse_args()
+    command: str | None = args.command
+
+    if command is None:
+        parser.print_help()
+        sys.exit(1)
+
+    if command == "generate":
+        root: Path = args.root
+        generate_command(root)
+
+
+def generate_command(root: Path) -> None:
+    resolved_root = root.resolve()
+    print(f"Scanning for StreamableBaseModel subclasses in: {resolved_root}")
+
+    models = find_streamable_models(resolved_root)
+
+    if not models:
+        print("No StreamableBaseModel subclasses found.")
+        return
+
+    print(f"Found {len(models)} model(s): {', '.join(m.__name__ for m in models)}")
+
+    code = generate_jmux_code(models)
+
+    output_path = Path(__file__).parent / "generated" / "__init__.py"
+    output_path.parent.mkdir(exist_ok=True)
+    output_path.write_text(code)
+
+    print(f"Generated JMux classes written to: {output_path}")
+
+
+if __name__ == "__main__":
+    main()

jmux/decoder.py

@@ -32,26 +32,40 @@ class StringEscapeDecoder:
     def __init__(self):
         self._buffer = ""
         self._string_escape = False
+        self._is_parsing_unicode = False
+        self._unicode_buffer = ""
 
     def push(self, ch: str) -> str | None:
+        if self._is_parsing_unicode:
+            self._unicode_buffer += ch
+            if len(self._unicode_buffer) == 4:
+                code_point = int(self._unicode_buffer, 16)
+                decoded_char = chr(code_point)
+                self._buffer += decoded_char
+                self._is_parsing_unicode = False
+                self._unicode_buffer = ""
+                return decoded_char
+            return None
+
         if self._string_escape:
             self._string_escape = False
             if ch == "u":
-                self.is_parsing_unicode = True
-                self.unicode_buffer = ""
-                return
+                self._is_parsing_unicode = True
+                self._unicode_buffer = ""
+                return None
             escaped_char = self.escape_map.get(ch, ch)
             self._buffer += escaped_char
             return escaped_char
 
         if ch == "\\":
             self._string_escape = True
+            return None
         else:
             self._buffer += ch
             return ch
 
     def is_terminating_quote(self, ch: str) -> bool:
-        if self._string_escape:
+        if self._string_escape or self._is_parsing_unicode:
             return False
         if ch == '"':
             return True
@@ -60,6 +74,8 @@ class StringEscapeDecoder:
     def reset(self) -> None:
         self._buffer = ""
         self._string_escape = False
+        self._is_parsing_unicode = False
+        self._unicode_buffer = ""
 
     @property
     def buffer(self) -> str:

jmux/demux.py

@@ -47,6 +47,8 @@ from jmux.types import (
     BOOLEAN_OPEN,
     COLON,
     COMMA,
+    EXPECT_KEY_IN_ROOT,
+    EXPECT_VALUE_IN_ARRAY,
     FLOAT_ALLOWED,
     INTERGER_ALLOWED,
     JSON_FALSE,
@@ -199,60 +201,74 @@ class JMux(ABC):
             ForbiddenTypeHintsError: If a type hint is not allowed.
             ObjectMissmatchedError: If the JMux class does not match the Pydantic model.
         """
-        for attr_name, type_alias in get_type_hints(cls).items():
-            jmux_main_type_set, jmux_subtype_set = extract_types_from_generic_alias(
-                type_alias
-            )
-
-            MaybePydanticType = get_type_hints(pydantic_model).get(attr_name, None)
-            if MaybePydanticType is None:
-                if NoneType in jmux_subtype_set:
-                    continue
-                else:
-                    raise MissingAttributeError(
-                        object_name=pydantic_model.__name__,
-                        attribute=attr_name,
-                    )
+        try:
+            for attr_name, type_alias in get_type_hints(cls).items():
+                jmux_main_type_set, jmux_subtype_set = extract_types_from_generic_alias(
+                    type_alias
+                )
 
-            pydantic_main_type_set, pydantic_subtype_set = (
-                extract_types_from_generic_alias(MaybePydanticType)
-            )
-            cls._assert_only_allowed_types(
-                jmux_main_type_set,
-                jmux_subtype_set,
-                pydantic_main_type_set,
-                pydantic_subtype_set,
-            )
-            cls._assert_correct_set_combinations(
-                jmux_main_type_set,
-                pydantic_main_type_set,
-                pydantic_subtype_set,
-            )
+                MaybePydanticType = get_type_hints(pydantic_model).get(attr_name, None)
+                if MaybePydanticType is None:
+                    if NoneType in jmux_subtype_set:
+                        continue
+                    else:
+                        raise MissingAttributeError(
+                            object_name=pydantic_model.__name__,
+                            attribute=attr_name,
+                        )
 
-            if StreamableValues in jmux_main_type_set:
-                cls._assert_is_allowed_streamable_values(
-                    jmux_subtype_set,
-                    pydantic_subtype_set,
-                    pydantic_main_type_set,
-                    pydantic_model,
-                    attr_name,
+                pydantic_main_type_set, pydantic_subtype_set = (
+                    extract_types_from_generic_alias(MaybePydanticType)
                 )
-            elif AwaitableValue in jmux_main_type_set:
-                cls._assert_is_allowed_awaitable_value(
+                cls._assert_only_allowed_types(
+                    jmux_main_type_set,
                     jmux_subtype_set,
-                    pydantic_subtype_set,
                     pydantic_main_type_set,
-                    pydantic_model,
-                    attr_name,
+                    pydantic_subtype_set,
                 )
-            else:
-                raise ObjectMissmatchedError(
-                    jmux_model=cls.__name__,
-                    pydantic_model=pydantic_model.__name__,
-                    attribute=attr_name,
-                    message="Unexpected main type on JMux",
+                cls._assert_correct_set_combinations(
+                    jmux_main_type_set,
+                    pydantic_main_type_set,
+                    pydantic_subtype_set,
                 )
 
+                if StreamableValues in jmux_main_type_set:
+                    cls._assert_is_allowed_streamable_values(
+                        jmux_subtype_set,
+                        pydantic_subtype_set,
+                        pydantic_main_type_set,
+                        pydantic_model,
+                        attr_name,
+                    )
+                elif AwaitableValue in jmux_main_type_set:
+                    cls._assert_is_allowed_awaitable_value(
+                        jmux_subtype_set,
+                        pydantic_subtype_set,
+                        pydantic_main_type_set,
+                        pydantic_model,
+                        attr_name,
+                    )
+                else:
+                    raise ObjectMissmatchedError(
+                        jmux_model=cls.__name__,
+                        pydantic_model=pydantic_model.__name__,
+                        attribute=attr_name,
+                        message="Unexpected main type on JMux",
+                    )
+        except TypeError as e:
+            raise ForbiddenTypeHintsError(message=f"Unexpected type hint: {e}") from e
+        except MissingAttributeError as e:
+            raise e
+        except ObjectMissmatchedError as e:
+            raise e
+        except Exception as e:
+            raise ObjectMissmatchedError(
+                jmux_model=cls.__name__,
+                pydantic_model=pydantic_model.__name__,
+                attribute=attr_name,
+                message=f"Unexpected main type on JMux: {e}",
+            )
+
     @classmethod
     def _assert_correct_set_combinations(
         cls,
@@ -481,18 +497,27 @@ class JMux(ABC):
             # CONTEXT: Root
             case M.ROOT:
                 match self._pda.state:
-                    case S.EXPECT_KEY:
+                    case _ if self._pda.state in EXPECT_KEY_IN_ROOT:
                         if ch in JSON_WHITESPACE:
                             pass
                         elif ch == '"':
                             self._pda.set_state(S.PARSING_KEY)
                             self._decoder.reset()
+                        elif ch in OBJECT_CLOSE:
+                            if self._pda.state is S.EXPECT_KEY_AFTER_COMMA:
+                                raise UnexpectedCharacterError(
+                                    ch,
+                                    self._pda.stack,
+                                    self._pda.state,
+                                    "Trailing comma in object is not allowed.",
+                                )
+                            await self._finalize()
                         else:
                             raise UnexpectedCharacterError(
                                 ch,
                                 self._pda.stack,
                                 self._pda.state,
-                                "Char needs to be '\"' or JSON whitespaces",
+                                "Char needs to be '\"', '}' or JSON whitespaces",
                             )
 
                     case S.PARSING_KEY:
@@ -537,6 +562,14 @@ class JMux(ABC):
                                     "Expected '[' or '\"' for 'StreamableValues'",
                                 )
                         elif ch in ARRAY_OPEN:
+                            if self._sink.current_sink_type == SinkType.AWAITABLE_VALUE:
+                                raise UnexpectedCharacterError(
+                                    ch,
+                                    self._pda.stack,
+                                    self._pda.state,
+                                    "Trying to parse 'array' but sink type is "
+                                    "'AwaitableValue'.",
+                                )
                             self._pda.set_state(S.EXPECT_VALUE)
                             self._pda.push(M.ARRAY)
                         else:
@@ -579,12 +612,12 @@ class JMux(ABC):
                             await self._parse_primitive()
                             await self._sink.close()
                             self._decoder.reset()
-                            if ch in JSON_WHITESPACE:
-                                self._pda.set_state(S.EXPECT_COMMA_OR_EOC)
-                            else:
-                                self._pda.set_state(S.EXPECT_KEY)
-                            if ch in OBJECT_CLOSE:
+                            if ch in COMMA:
+                                self._pda.set_state(S.EXPECT_KEY_AFTER_COMMA)
+                            elif ch in OBJECT_CLOSE:
                                 await self._finalize()
+                            else:
+                                self._pda.set_state(S.EXPECT_COMMA_OR_EOC)
                         else:
                             self._assert_primitive_character_allowed_in_state(ch)
                             self._decoder.push(ch)
@@ -593,7 +626,7 @@ class JMux(ABC):
                         if ch in JSON_WHITESPACE:
                             pass
                         elif ch in COMMA:
-                            self._pda.set_state(S.EXPECT_KEY)
+                            self._pda.set_state(S.EXPECT_KEY_AFTER_COMMA)
                         elif ch in OBJECT_CLOSE:
                             await self._finalize()
                         else:
@@ -622,12 +655,19 @@ class JMux(ABC):
                     )
 
                 match self._pda.state:
-                    case S.EXPECT_VALUE:
+                    case _ if self._pda.state in EXPECT_VALUE_IN_ARRAY:
                         if ch in JSON_WHITESPACE:
                             pass
                         elif await self._handle_common__expect_value(ch):
                             pass
                         elif ch in ARRAY_CLOSE:
+                            if self._pda.state is S.EXPECT_VALUE_AFTER_COMMA:
+                                raise UnexpectedCharacterError(
+                                    ch,
+                                    self._pda.stack,
+                                    self._pda.state,
+                                    "Trailing comma in array is not allowed.",
+                                )
                             await self._close_context(S.EXPECT_COMMA_OR_EOC)
                         else:
                             raise UnexpectedCharacterError(
@@ -670,9 +710,11 @@ class JMux(ABC):
                             await self._parse_primitive()
                             self._decoder.reset()
                             if ch in COMMA:
-                                self._pda.set_state(S.EXPECT_VALUE)
+                                self._pda.set_state(S.EXPECT_VALUE_AFTER_COMMA)
                             elif ch in ARRAY_CLOSE:
                                 await self._close_context(S.EXPECT_COMMA_OR_EOC)
+                            else:
+                                self._pda.set_state(S.EXPECT_COMMA_OR_EOC)
                         else:
                             self._assert_primitive_character_allowed_in_state(ch)
                             self._decoder.push(ch)
@@ -681,7 +723,7 @@ class JMux(ABC):
                         if ch in JSON_WHITESPACE:
                             pass
                         elif ch in COMMA:
-                            self._pda.set_state(S.EXPECT_VALUE)
+                            self._pda.set_state(S.EXPECT_VALUE_AFTER_COMMA)
                         elif ch in ARRAY_CLOSE:
                             await self._close_context(S.EXPECT_COMMA_OR_EOC)
                         else:

jmux/generator.py

@@ -0,0 +1,381 @@
+from __future__ import annotations
+
+import ast
+import importlib.util
+import sys
+from enum import Enum
+from pathlib import Path
+from types import ModuleType, NoneType, UnionType
+from typing import Annotated, Any, Union, get_args, get_origin, get_type_hints
+
+from pydantic import BaseModel
+
+from jmux.base import StreamableBaseModel, Streamed
+
+
+def extract_models_from_source(
+    source: str,
+    module_name: str = "__jmux_extract__",
+) -> list[type[StreamableBaseModel]]:
+    if not _source_imports_streamable_base_model(source):
+        return []
+
+    module = _exec_source_as_module(source, module_name)
+    if module is None:
+        return []
+
+    return _extract_models_from_module(module)
+
+
+def find_streamable_models(root_path: Path) -> list[type[StreamableBaseModel]]:
+    models: list[type[StreamableBaseModel]] = []
+    root_path = root_path.resolve()
+
+    for py_file in root_path.rglob("*.py"):
+        source = _read_file_safe(py_file)
+        if source is None:
+            continue
+
+        if not _source_imports_streamable_base_model(source):
+            continue
+
+        module = _import_module_from_path(py_file, root_path)
+        if module is None:
+            continue
+
+        models.extend(_extract_models_from_module(module))
+
+    return models
+
+
+def _extract_models_from_module(
+    module: ModuleType,
+) -> list[type[StreamableBaseModel]]:
+    models: list[type[StreamableBaseModel]] = []
+    for name in dir(module):
+        obj = getattr(module, name)
+        if (
+            isinstance(obj, type)
+            and issubclass(obj, StreamableBaseModel)
+            and obj is not StreamableBaseModel
+        ):
+            models.append(obj)
+    return models
+
+
+def _read_file_safe(py_file: Path) -> str | None:
+    try:
+        return py_file.read_text()
+    except (OSError, UnicodeDecodeError):
+        return None
+
+
+def _source_imports_streamable_base_model(source: str) -> bool:
+    try:
+        tree = ast.parse(source)
+    except SyntaxError:
+        return False
+
+    for node in ast.walk(tree):
+        if isinstance(node, ast.ImportFrom):
+            if node.module and "jmux" in node.module:
+                for alias in node.names:
+                    if alias.name == "StreamableBaseModel":
+                        return True
+        elif isinstance(node, ast.Import):
+            for alias in node.names:
+                if "jmux" in alias.name:
+                    return True
+    return False
+
+
+def _exec_source_as_module(source: str, module_name: str) -> ModuleType | None:
+    try:
+        module = ModuleType(module_name)
+        module.__dict__["__builtins__"] = __builtins__
+        sys.modules[module_name] = module
+        exec(compile(source, f"<{module_name}>", "exec"), module.__dict__)
+        return module
+    except Exception:
+        if module_name in sys.modules:
+            del sys.modules[module_name]
+        return None
+
+
+def _import_module_from_path(py_file: Path, root_path: Path) -> ModuleType | None:
+    try:
+        relative = py_file.relative_to(root_path)
+        module_name = str(relative.with_suffix("")).replace("/", ".").replace("\\", ".")
+
+        if root_path not in [Path(p) for p in sys.path]:
+            sys.path.insert(0, str(root_path))
+
+        spec = importlib.util.spec_from_file_location(module_name, py_file)
+        if spec is None or spec.loader is None:
+            return None
+
+        module = importlib.util.module_from_spec(spec)
+        sys.modules[module_name] = module
+        spec.loader.exec_module(module)
+        return module
+    except Exception:
+        return None
+
+
+def _is_streamed_marker(obj: Any) -> bool:
+    if obj is Streamed:
+        return True
+    if isinstance(obj, type) and obj.__name__ == "Streamed":
+        return True
+    return False
+
+
+def _get_resolved_annotations(model: type[StreamableBaseModel]) -> dict[str, Any]:
+    try:
+        globalns = getattr(sys.modules.get(model.__module__, None), "__dict__", {})
+        return get_type_hints(model, globalns=globalns, include_extras=True)
+    except Exception:
+        return model.__annotations__
+
+
+def generate_jmux_code(models: list[type[StreamableBaseModel]]) -> str:
+    lines = [
+        "from enum import Enum",
+        "from typing import Union",
+        "",
+        "from jmux.awaitable import AwaitableValue, StreamableValues",
+        "from jmux.demux import JMux",
+        "",
+    ]
+
+    enum_imports = _collect_enum_imports(models)
+    if enum_imports:
+        lines.extend(enum_imports)
+        lines.append("")
+
+    nested_models = _collect_nested_models(models)
+    all_models = _topological_sort(models, nested_models)
+
+    for model in all_models:
+        class_lines = _generate_class(model)
+        lines.extend(class_lines)
+        lines.append("")
+
+    return "\n".join(lines)
+
+
+def _collect_enum_imports(models: list[type[StreamableBaseModel]]) -> list[str]:
+    enums: set[type[Enum]] = set()
+    for model in models:
+        resolved = _get_resolved_annotations(model)
+        for field_name in model.model_fields:
+            annotation = resolved.get(field_name)
+            if annotation is None:
+                continue
+            _collect_enums_from_type(annotation, enums)
+
+    imports = []
+    for enum_type in enums:
+        module = enum_type.__module__
+        name = enum_type.__name__
+        imports.append(f"from {module} import {name}")
+    return sorted(imports)
+
+
+def _collect_enums_from_type(annotation: Any, enums: set[type[Enum]]) -> None:
+    origin = get_origin(annotation)
+    args = get_args(annotation)
+
+    if origin is Annotated:
+        if args:
+            _collect_enums_from_type(args[0], enums)
+        return
+
+    if isinstance(annotation, type) and issubclass(annotation, Enum):
+        enums.add(annotation)
+        return
+
+    if origin is list and args:
+        _collect_enums_from_type(args[0], enums)
+
+
+def _collect_nested_models(
+    models: list[type[StreamableBaseModel]],
+) -> set[type[StreamableBaseModel]]:
+    nested: set[type[StreamableBaseModel]] = set()
+    for model in models:
+        resolved = _get_resolved_annotations(model)
+        for field_name in model.model_fields:
+            annotation = resolved.get(field_name)
+            if annotation is None:
+                continue
+            _collect_nested_from_type(annotation, nested)
+    return nested
+
+
+def _collect_nested_from_type(
+    annotation: Any, nested: set[type[StreamableBaseModel]]
+) -> None:
+    origin = get_origin(annotation)
+    args = get_args(annotation)
+
+    if origin is Annotated:
+        if args:
+            _collect_nested_from_type(args[0], nested)
+        return
+
+    if (
+        isinstance(annotation, type)
+        and issubclass(annotation, BaseModel)
+        and issubclass(annotation, StreamableBaseModel)
+        and annotation is not StreamableBaseModel
+    ):
+        nested.add(annotation)
+        return
+
+    if origin is list and args:
+        _collect_nested_from_type(args[0], nested)
+
+
+def _topological_sort(
+    models: list[type[StreamableBaseModel]],
+    nested: set[type[StreamableBaseModel]],
+) -> list[type[StreamableBaseModel]]:
+    all_models = set(models) | nested
+    result: list[type[StreamableBaseModel]] = []
+    visited: set[type[StreamableBaseModel]] = set()
+
+    def visit(model: type[StreamableBaseModel]) -> None:
+        if model in visited:
+            return
+        visited.add(model)
+        resolved = _get_resolved_annotations(model)
+        for field_name in model.model_fields:
+            annotation = resolved.get(field_name)
+            if annotation is None:
+                continue
+            dep = _get_nested_model_dependency(annotation)
+            if dep and dep in all_models:
+                visit(dep)
+        result.append(model)
+
+    for model in all_models:
+        visit(model)
+
+    return result
+
+
+def _get_nested_model_dependency(annotation: Any) -> type[StreamableBaseModel] | None:
+    origin = get_origin(annotation)
+    args = get_args(annotation)
+
+    if origin is Annotated:
+        if args:
+            return _get_nested_model_dependency(args[0])
+        return None
+
+    if (
+        isinstance(annotation, type)
+        and issubclass(annotation, BaseModel)
+        and issubclass(annotation, StreamableBaseModel)
+        and annotation is not StreamableBaseModel
+    ):
+        return annotation
+
+    if origin is list and args:
+        return _get_nested_model_dependency(args[0])
+
+    return None
+
+
+def _generate_class(model: type[StreamableBaseModel]) -> list[str]:
+    class_name = f"{model.__name__}JMux"
+    lines = [f"class {class_name}(JMux):"]
+
+    resolved = _get_resolved_annotations(model)
+    has_fields = False
+    for field_name in model.model_fields:
+        annotation = resolved.get(field_name)
+        if annotation is None:
+            continue
+        jmux_type = get_jmux_type(annotation)
+        lines.append(f"    {field_name}: {jmux_type}")
+        has_fields = True
+
+    if not has_fields:
+        lines.append("    pass")
+
+    return lines
+
+
+def get_jmux_type(annotation: Any) -> str:
+    origin = get_origin(annotation)
+    args = get_args(annotation)
+
+    if origin is Annotated:
+        inner_type = args[0] if args else None
+        metadata = args[1:] if len(args) > 1 else ()
+        has_streamed = any(_is_streamed_marker(m) for m in metadata)
+
+        if has_streamed and inner_type is str:
+            return "StreamableValues[str]"
+
+        return get_jmux_type(inner_type)
+
+    if origin is list:
+        inner = args[0] if args else Any
+        inner_jmux = _get_inner_type_str(inner)
+        return f"StreamableValues[{inner_jmux}]"
+
+    if annotation is str:
+        return "AwaitableValue[str]"
+
+    if annotation is int:
+        return "AwaitableValue[int]"
+
+    if annotation is float:
+        return "AwaitableValue[float]"
+
+    if annotation is bool:
+        return "AwaitableValue[bool]"
+
+    if annotation is NoneType or annotation is None:
+        return "AwaitableValue[None]"
+
+    if isinstance(annotation, type) and issubclass(annotation, Enum):
+        return f"AwaitableValue[{annotation.__name__}]"
+
+    if isinstance(annotation, type) and issubclass(annotation, BaseModel):
+        if issubclass(annotation, StreamableBaseModel):
+            return f"AwaitableValue[{annotation.__name__}JMux]"
+        return f"AwaitableValue[{annotation.__name__}]"
+
+    type_origin = get_origin(annotation)
+    if type_origin is UnionType or type_origin is Union:
+        non_none_args = [a for a in args if a is not NoneType and a is not None]
+        if len(non_none_args) == 1:
+            inner = _get_inner_type_str(non_none_args[0])
+            return f"AwaitableValue[{inner} | None]"
+        return f"AwaitableValue[{annotation}]"
+
+    return f"AwaitableValue[{annotation}]"
+
+
+def _get_inner_type_str(annotation: Any) -> str:
+    if annotation is str:
+        return "str"
+    if annotation is int:
+        return "int"
+    if annotation is float:
+        return "float"
+    if annotation is bool:
+        return "bool"
+    if annotation is NoneType or annotation is None:
+        return "None"
+    if isinstance(annotation, type) and issubclass(annotation, Enum):
+        return annotation.__name__
+    if isinstance(annotation, type) and issubclass(annotation, BaseModel):
+        if issubclass(annotation, StreamableBaseModel):
+            return f"{annotation.__name__}JMux"
+        return annotation.__name__
+    return str(annotation)

jmux/types.py

@@ -9,8 +9,10 @@ class State(Enum):
     ERROR = "error"
     # expect
     EXPECT_KEY = "expect_key"
+    EXPECT_KEY_AFTER_COMMA = "expect_key_after_comma"
     EXPECT_COLON = "expect_colon"
     EXPECT_VALUE = "expect_value"
+    EXPECT_VALUE_AFTER_COMMA = "expect_value_after_comma"
     EXPECT_COMMA_OR_EOC = "expect_comma_or_eoc"
     # parsing
     PARSING_KEY = "parsing_key"
@@ -29,6 +31,9 @@ PARSING_PRIMITIVE_STATES: Set[State] = {
     State.PARSING_NULL,
 }
 
+EXPECT_KEY_IN_ROOT = {State.EXPECT_KEY, State.EXPECT_KEY_AFTER_COMMA}
+EXPECT_VALUE_IN_ARRAY = {State.EXPECT_VALUE, State.EXPECT_VALUE_AFTER_COMMA}
+
 
 class Mode(Enum):
     ROOT = "$"

{jmux-0.0.5.dist-info → jmux-0.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jmux
-Version: 0.0.5
+Version: 0.1.0
 Summary: JMux: A Python package for demultiplexing a JSON string into multiple awaitable variables.
 Author-email: "Johannes A.I. Unruh" <johannes@unruh.ai>
 License: MIT License
@@ -40,10 +40,12 @@ Requires-Dist: pydantic>=2.0.0
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-anyio; extra == "test"
+Requires-Dist: trio; extra == "test"
 Provides-Extra: dev
 Requires-Dist: ruff; extra == "dev"
 Requires-Dist: pytest; extra == "dev"
 Requires-Dist: pytest-anyio; extra == "dev"
+Requires-Dist: trio; extra == "dev"
 Requires-Dist: uv; extra == "dev"
 Requires-Dist: build; extra == "dev"
 Requires-Dist: twine; extra == "dev"
@@ -61,8 +63,10 @@ This package is inspired by `Snapshot Streaming` mentioned in the [`WWDC25: Meet
 
 ## Features
 
-- **Asynchronous by Design**: Built on top of `asyncio`, JMux is perfect for modern, high-performance Python applications.
+- **Asynchronous by Design**: Built on top of `anyio`, JMux supports both `asyncio` and `trio` backends, making it perfect for modern, high-performance Python applications.
+- **Python 3.10+**: Supports Python 3.10 and newer versions.
 - **Pydantic Integration**: Validate your `JMux` classes against Pydantic models to ensure type safety and consistency.
+- **Code Generation**: Automatically generate JMux classes from `StreamableBaseModel` subclasses using the `jmux generate` CLI command.
 - **Awaitable and Streamable Sinks**: Use `AwaitableValue` for single values and `StreamableValues` for streams of values.
 - **Robust Error Handling**: JMux provides a comprehensive set of exceptions to handle parsing errors and other issues.
 - **Lightweight**: JMux has only a few external dependencies, making it easy to integrate into any project.
@@ -75,21 +79,155 @@ You can install JMux from PyPI using pip:
 pip install jmux
 ```
 
+## Code Generation
+
+JMux provides a CLI tool to automatically generate JMux classes from Pydantic models. Instead of manually writing both a Pydantic model and a corresponding JMux class, you can define your models using `StreamableBaseModel` and let JMux generate the demultiplexer classes for you.
+
+### Defining Models with StreamableBaseModel
+
+Use `StreamableBaseModel` as your base class instead of `pydantic.BaseModel`:
+
+```python
+from typing import Annotated
+from jmux import StreamableBaseModel, Streamed
+
+class LlmResponse(StreamableBaseModel):
+    thought: str
+    tool_code: Annotated[str, Streamed]
+    tags: list[str]
+```
+
+### Type Mappings
+
+The generator converts your model fields to JMux types as follows:
+
+| Model Field Type              | Generated JMux Type               |
+| ----------------------------- | --------------------------------- |
+| `str`, `int`, `float`, `bool` | `AwaitableValue[T]`               |
+| `Enum`                        | `AwaitableValue[EnumType]`        |
+| `T \| None`                   | `AwaitableValue[T \| None]`       |
+| `list[T]`                     | `StreamableValues[T]`             |
+| `Annotated[str, Streamed]`    | `StreamableValues[str]`           |
+| Nested `StreamableBaseModel`  | `AwaitableValue[NestedModelJMux]` |
+
+The `Streamed` marker is useful when you want to stream a string field character-by-character (e.g., for real-time display of LLM output) rather than awaiting the complete value.
+
+### Using the CLI
+
+Run the `jmux generate` command to scan your codebase and generate JMux classes:
+
+```bash
+jmux generate --root <directory>
+```
+
+This will:
+
+1. Recursively scan `<directory>` for Python files containing `StreamableBaseModel` subclasses
+2. Generate corresponding JMux classes with the suffix `JMux` (e.g., `LlmResponse` → `LlmResponseJMux`)
+3. Write the generated code to `src/jmux/generated/__init__.py`
+
+### Example
+
+Given this model:
+
+```python
+from typing import Annotated
+from jmux import StreamableBaseModel, Streamed
+
+class LlmResponse(StreamableBaseModel):
+    thought: str
+    tool_code: Annotated[str, Streamed]
+```
+
+Running `jmux generate` produces:
+
+```python
+from jmux.awaitable import AwaitableValue, StreamableValues
+from jmux.demux import JMux
+
+class LlmResponseJMux(JMux):
+    thought: AwaitableValue[str]
+    tool_code: StreamableValues[str]
+```
+
+You can then import and use the generated class:
+
+```python
+from jmux.generated import LlmResponseJMux
+
+jmux_instance = LlmResponseJMux()
+```
+
 ## Usage with LLMs (e.g., `litellm`)
 
 The primary use case for `jmux` is to process streaming JSON responses from LLMs. This allows you to react to parts of the data as it arrives, rather than waiting for the entire JSON object to be transmitted. While this should be obvious, I should mention, that **the order in which the pydantic model defines the properties, defines which stream is filled first**.
 
-Here’s a conceptual example of how you might integrate `jmux` with an LLM call, such as one made with `litellm`:
+### Using Code Generation (Recommended)
+
+The easiest way to use JMux with LLMs is to define your models using `StreamableBaseModel` and generate JMux classes automatically:
+
+```python
+from typing import Annotated
+from jmux import StreamableBaseModel, Streamed
+
+class LlmResponse(StreamableBaseModel): # Use `StreamableBaseModel` so that the CLI can find the `pydantic` models to parse
+    thought: str
+    tool_code: Annotated[str, Streamed]
+```
+
+Then run `jmux generate --root .` to generate the `LlmResponseJMux` class. You can then use it directly:
+
+```python
+import anyio
+from jmux.generated import LlmResponseJMux
+
+async def mock_llm_stream():
+    json_stream = '{"thought": "I need to write some code.", "tool_code": "print(\'Hello, World!\')"}'
+    for char in json_stream:
+        yield char
+        await anyio.sleep(0.01)
+
+async def process_llm_response():
+    jmux_instance = LlmResponseJMux()
+
+    async def feed_stream():
+        async for chunk in mock_llm_stream():
+            await jmux_instance.feed_chunks(chunk)
+
+    async def consume_thought():
+        thought = await jmux_instance.thought
+        print(f"LLM's thought received: '{thought}'")
+
+    async def consume_tool_code():
+        print("Receiving tool code...")
+        full_code = ""
+        async for code_fragment in jmux_instance.tool_code:
+            full_code += code_fragment
+            print(f"  -> Received fragment: {code_fragment}")
+        print(f"Full tool code received: {full_code}")
+
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(feed_stream)
+        tg.start_soon(consume_thought)
+        tg.start_soon(consume_tool_code)
+
+if __name__ == "__main__":
+    anyio.run(process_llm_response, backend="asyncio")
+```
+
+### Manual Approach
+
+If you prefer more control, you can manually define both the Pydantic model and the JMux class:
 
 ```python
-import asyncio
+import anyio
 from pydantic import BaseModel
 from jmux import JMux, AwaitableValue, StreamableValues
 # litellm is used conceptually here
 # from litellm import acompletion
 
 # 1. Define the Pydantic model for the expected JSON response
-class LlmResponse(BaseModel):
+class LlmResponse(BaseModel): # No need to use `StreamableBaseModel` here, since it is only used for detection purposes
     thought: str # **This property is filled first**
     tool_code: str
 
@@ -106,7 +244,7 @@ async def mock_llm_stream():
     json_stream = '{"thought": "I need to write some code.", "tool_code": "print(\'Hello, World!\')"}'
     for char in json_stream:
         yield char
-        await asyncio.sleep(0.01) # Simulate network latency
+        await anyio.sleep(0.01) # Simulate network latency
 
 # Main function to orchestrate the call and processing
 async def process_llm_response():
@@ -132,15 +270,16 @@ async def process_llm_response():
             print(f"  -> Received fragment: {code_fragment}")
         print(f"Full tool code received: {full_code}")
 
-    # Run all tasks concurrently
-    await asyncio.gather(
-        feed_stream(),
-        consume_thought(),
-        consume_tool_code()
-    )
+    # Run all tasks concurrently using anyio task group
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(feed_stream)
+        tg.start_soon(consume_thought)
+        tg.start_soon(consume_tool_code)
 
+# Run with asyncio backend
 if __name__ == "__main__":
-    asyncio.run(process_llm_response())
+    anyio.run(process_llm_response, backend="asyncio")
+    # Or use trio: anyio.run(process_llm_response, backend="trio")
 ```
 
 ## Example Implementation
@@ -212,7 +351,7 @@ You can either `await awaitable_llm_result` if you need the full result, or use
 Here is a simple example of how to use JMux to parse a JSON stream:
 
 ```python
-import asyncio
+import anyio
 from enum import Enum
 from types import NoneType
 from pydantic import BaseModel
@@ -297,10 +436,13 @@ async def main():
         nested_key_str = await key_nested.key_str
         print(f"nested_key_str: {nested_key_str}")
 
-    await asyncio.gather(produce(), consume())
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(produce)
+        tg.start_soon(consume)
 
 if __name__ == "__main__":
-    asyncio.run(main())
+    anyio.run(main, backend="asyncio")
+    # Or use trio: anyio.run(main, backend="trio")
 ```
 
 ## API Reference
@@ -349,13 +491,35 @@ Additionally the following type is supported without being wrapped into `list`:
 
 This allows you to fully stream strings directly to a sink.
 
-## License
+### Class `jmux.StreamableBaseModel`
 
-This project is licensed under the terms of the MIT license. See the [LICENSE](LICENSE) file for details.
+A Pydantic `BaseModel` subclass used for defining models that can be automatically converted to JMux classes via the `jmux generate` CLI command.
+
+```python
+from jmux import StreamableBaseModel
 
-## Planned Improvements
+class MyModel(StreamableBaseModel):
+    name: str
+    age: int
+```
+
+### Class `jmux.Streamed`
+
+A marker class used with `typing.Annotated` to indicate that a string field should be streamed character-by-character rather than awaited as a complete value.
 
-- Add support for older Python versions
+```python
+from typing import Annotated
+from jmux import StreamableBaseModel, Streamed
+
+class MyModel(StreamableBaseModel):
+    content: Annotated[str, Streamed]
+```
+
+When generating JMux classes, fields annotated with `Streamed` will be converted to `StreamableValues[str]` instead of `AwaitableValue[str]`.
+
+## License
+
+This project is licensed under the terms of the MIT license. See the [LICENSE](LICENSE) file for details.
 
 ## Contributions
 

jmux-0.1.0.dist-info/RECORD

@@ -0,0 +1,17 @@
+jmux/__init__.py,sha256=ji94yf4_4CyqNfLmJWoJn9WXWJyiRhJGmUqFOiHYdFc,256
+jmux/awaitable.py,sha256=MV7kltI_3D-OzNrTSfiRi8xmARl0Wu8db0QhQ4-zom0,9038
+jmux/base.py,sha256=v1T6_lYicX05a_KbIy5v3ZbnuR_PQmefN0ruaP5DOnM,107
+jmux/cli.py,sha256=wfAcZKMwY_e7Zfd3XecTVgCURC43PEM77k4eZ6_dgIE,1611
+jmux/decoder.py,sha256=YoCkBKposMx1PA78BIVawgFw96QsqTZaOW4gNCSkI7I,2197
+jmux/demux.py,sha256=R2vazLOtfh4zvQKoj3_92RgdqG-zaG3CkQMBz72fAIE,38800
+jmux/error.py,sha256=VZJYivt8RPfjcF2bs-T7_UkH3dVA3xH-xGbZggQV14k,4665
+jmux/generator.py,sha256=clKTZHh-VNsRJarEa8MBT3-zXYe-8CZnI6QbJbRPo-M,11408
+jmux/helpers.py,sha256=zOlw1Yk7-sdKAeasswRRcuUOTEBAUbymoAGwBTMaOjg,2902
+jmux/pda.py,sha256=19joQd0DD5OAmwRpp2jVVbtiFXnjv5P_1mZm87-QOeY,922
+jmux/types.py,sha256=a5Xyx_8A3gsZkscII1_7E6oyu4VQ-oI4_osTnPLi_xs,1649
+jmux-0.1.0.dist-info/licenses/LICENSE,sha256=y0qnwaAe4bEqzNPyq4M_VZA2I2mQly8MawajyZhqw0k,1169
+jmux-0.1.0.dist-info/METADATA,sha256=DbM4_Qf6cGXxo5vSiZ7JOkuVadsZ1s62lsUwPL_UeyE,19004
+jmux-0.1.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+jmux-0.1.0.dist-info/entry_points.txt,sha256=8YYzNePAj4-3v1s6IJ_esFhgg9TZZL2fRYXMbSgmAho,39
+jmux-0.1.0.dist-info/top_level.txt,sha256=TF2N6kHqLghfOkCiNlCueMDX4l5rPn_5MSPNtYrS1-o,5
+jmux-0.1.0.dist-info/RECORD,,

jmux-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+jmux = jmux.cli:main

jmux-0.0.5.dist-info/RECORD

@@ -1,13 +0,0 @@
-jmux/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jmux/awaitable.py,sha256=otX4vbWzVmIwxvMDWOtaIvS9do8YJtbls9xFjwH9Yw0,8534
-jmux/decoder.py,sha256=Y6KVryRDLvGV5nBsneXpTvC0WUGhR5Z89Dvqz4HMAgg,1562
-jmux/demux.py,sha256=OuUaNwvKI6WJxdzBzbBgx0yPq1kn4jSttuB5lvwzztk,36422
-jmux/error.py,sha256=VZJYivt8RPfjcF2bs-T7_UkH3dVA3xH-xGbZggQV14k,4665
-jmux/helpers.py,sha256=zOlw1Yk7-sdKAeasswRRcuUOTEBAUbymoAGwBTMaOjg,2902
-jmux/pda.py,sha256=19joQd0DD5OAmwRpp2jVVbtiFXnjv5P_1mZm87-QOeY,922
-jmux/types.py,sha256=CJhFS9RVgR0cDBNJR8ROAFnxzG4YTYpNZ90hyD2SxsY,1389
-jmux-0.0.5.dist-info/licenses/LICENSE,sha256=y0qnwaAe4bEqzNPyq4M_VZA2I2mQly8MawajyZhqw0k,1169
-jmux-0.0.5.dist-info/METADATA,sha256=31D2JB7C5kqeislSLEg3WRkyGXO5cx3mY34SAJd1ntU,13330
-jmux-0.0.5.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-jmux-0.0.5.dist-info/top_level.txt,sha256=TF2N6kHqLghfOkCiNlCueMDX4l5rPn_5MSPNtYrS1-o,5
-jmux-0.0.5.dist-info/RECORD,,