fluent-codegen 0.1.0__tar.gz → 0.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: fluent-codegen
3
- Version: 0.1.0
3
+ Version: 0.2
4
4
  Summary: A Python library for generating Python code via AST construction.
5
5
  Author-email: Luke Plant <luke@lukeplant.me.uk>
6
6
  License-Expression: Apache-2.0
@@ -76,7 +76,7 @@ fluent method-chaining for expressions:
76
76
  func, _ = module.create_function("fizzbuzz", args=["n"])
77
77
 
78
78
  # 2. A Name reference to the "n" parameter (Function *is* a Scope)
79
- n = codegen.Name("n", func)
79
+ n = func.name("n")
80
80
 
81
81
  # 3. Build an if / elif / else chain
82
82
  if_stmt = func.body.create_if()
@@ -94,7 +94,7 @@ fluent method-chaining for expressions:
94
94
  branch.create_return(codegen.String("Buzz"))
95
95
 
96
96
  # else: return str(n)
97
- if_stmt.else_block.create_return(codegen.function_call("str", [n], {}, func))
97
+ if_stmt.else_block.create_return(module.scope.name("str").call([n]))
98
98
 
99
99
  # 4. Inspect the generated source
100
100
  print(module.as_python_source())
@@ -54,7 +54,7 @@ fluent method-chaining for expressions:
54
54
  func, _ = module.create_function("fizzbuzz", args=["n"])
55
55
 
56
56
  # 2. A Name reference to the "n" parameter (Function *is* a Scope)
57
- n = codegen.Name("n", func)
57
+ n = func.name("n")
58
58
 
59
59
  # 3. Build an if / elif / else chain
60
60
  if_stmt = func.body.create_if()
@@ -72,7 +72,7 @@ fluent method-chaining for expressions:
72
72
  branch.create_return(codegen.String("Buzz"))
73
73
 
74
74
  # else: return str(n)
75
- if_stmt.else_block.create_return(codegen.function_call("str", [n], {}, func))
75
+ if_stmt.else_block.create_return(module.scope.name("str").call([n]))
76
76
 
77
77
  # 4. Inspect the generated source
78
78
  print(module.as_python_source())
@@ -24,8 +24,7 @@ classifiers = [
24
24
  "Topic :: Software Development :: Code Generators",
25
25
  "Topic :: Software Development :: Libraries :: Python Modules",
26
26
  ]
27
- dynamic = ["version"]
28
-
27
+ version = "0.2"
29
28
  dependencies = []
30
29
 
31
30
  [project.urls]
@@ -48,8 +47,6 @@ include-package-data = false
48
47
  where = ["src"]
49
48
  namespaces = false
50
49
 
51
- [tool.setuptools.dynamic]
52
- version = {attr = "fluent_codegen.__version__"}
53
50
 
54
51
  [tool.ruff]
55
52
  line-length = 120
@@ -11,6 +11,7 @@ from typing import TypedDict
11
11
 
12
12
  Add = ast.Add
13
13
  And = ast.And
14
+ Assert = ast.Assert
14
15
  Assign = ast.Assign
15
16
  AnnAssign = ast.AnnAssign
16
17
  BoolOp = ast.BoolOp
@@ -77,6 +78,12 @@ alias = ast.alias
77
78
  # It's hard to get something sensible we can put for line/col numbers so we put arbitrary values.
78
79
 
79
80
 
81
+ try:
82
+ _Unparser = ast._Unparser # type: ignore
83
+ except AttributeError: # pragma: no cover
84
+ from _ast_unparse import Unparser as _Unparser # pragma: no cover # type: ignore
85
+
86
+
80
87
  class DefaultAstArgs(TypedDict):
81
88
  lineno: int
82
89
  col_offset: int
@@ -91,3 +98,31 @@ DEFAULT_AST_ARGS_ARGUMENTS: dict[str, object] = dict()
91
98
 
92
99
  def subscript_slice_object[T](value: T) -> T:
93
100
  return value
101
+
102
+
103
+ class CommentNode(ast.stmt):
104
+ """Custom AST statement node representing a comment.
105
+
106
+ This is not a standard Python AST node. It is ignored by ``compile()``
107
+ (callers must strip it first, which ``as_ast()`` does automatically),
108
+ but is rendered by :func:`unparse_with_comments`.
109
+ """
110
+
111
+ _fields = ("text",)
112
+
113
+ def __init__(self, text: str, **kwargs: int) -> None:
114
+ self.text = text
115
+ super().__init__(**kwargs) # type: ignore[reportCallIssue]
116
+
117
+
118
+ class _CommentUnparser(_Unparser): # type: ignore[reportAttributeAccessIssue]
119
+ """An unparser that knows how to render :class:`CommentNode`."""
120
+
121
+ def visit_CommentNode(self, node: CommentNode) -> None:
122
+ self.fill("# " + node.text) # type: ignore[reportAttributeAccessIssue]
123
+
124
+
125
+ def unparse_with_comments(node: ast.AST) -> str:
126
+ """Like :func:`ast.unparse`, but also renders :class:`CommentNode` nodes."""
127
+ unparser = _CommentUnparser()
128
+ return unparser.visit(node) # type: ignore[reportUnknownMemberType,reportUnknownVariableType]
@@ -8,6 +8,7 @@ import builtins
8
8
  import enum
9
9
  import keyword
10
10
  import re
11
+ import textwrap
11
12
  from abc import ABC, abstractmethod
12
13
  from collections.abc import Callable, Sequence
13
14
  from dataclasses import dataclass
@@ -19,9 +20,14 @@ from .ast_compat import (
19
20
  DEFAULT_AST_ARGS_ADD,
20
21
  DEFAULT_AST_ARGS_ARGUMENTS,
21
22
  DEFAULT_AST_ARGS_MODULE,
23
+ CommentNode,
24
+ unparse_with_comments,
22
25
  )
23
26
  from .utils import allowable_keyword_arg_name, allowable_name
24
27
 
28
+ #: Type alias for comment strings stored in :attr:`Block.statements`.
29
+ Comment = str
30
+
25
31
  # This module provides simple utilities for building up Python source code.
26
32
  # The design originally came from fluent-compiler, so had the following aims
27
33
  # and constraints:
@@ -61,15 +67,6 @@ from .utils import allowable_keyword_arg_name, allowable_name
61
67
  # which have similar aims.
62
68
 
63
69
 
64
- PROPERTY_TYPE = "PROPERTY_TYPE"
65
- PROPERTY_RETURN_TYPE = "PROPERTY_RETURN_TYPE"
66
- # UNKNOWN_TYPE is just an alias for `object` for clarity.
67
- UNKNOWN_TYPE: type = object
68
- # It is important for our usage of it that UNKNOWN_TYPE is a `type`,
69
- # and the most general `type`.
70
- assert isinstance(UNKNOWN_TYPE, type)
71
-
72
-
73
70
  SENSITIVE_FUNCTIONS = {
74
71
  # builtin functions that we should never be calling from our code
75
72
  # generation. This is a defense-in-depth mechansim to stop our code
@@ -103,15 +100,15 @@ class CodeGenAst(ABC):
103
100
  """
104
101
 
105
102
  @abstractmethod
106
- def as_ast(self) -> py_ast.AST: ...
103
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.AST: ...
107
104
 
108
105
  child_elements: ClassVar[list[str]]
109
106
 
110
107
  def as_python_source(self) -> str:
111
108
  """Return the Python source code for this AST node."""
112
- node = self.as_ast()
109
+ node = self.as_ast(include_comments=True)
113
110
  py_ast.fix_missing_locations(node)
114
- return py_ast.unparse(node)
111
+ return unparse_with_comments(node)
115
112
 
116
113
 
117
114
  class CodeGenAstList(ABC):
@@ -121,15 +118,15 @@ class CodeGenAstList(ABC):
121
118
  """
122
119
 
123
120
  @abstractmethod
124
- def as_ast_list(self, allow_empty: bool = True) -> list[py_ast.stmt]: ...
121
+ def as_ast_list(self, allow_empty: bool = True, *, include_comments: bool = False) -> list[py_ast.stmt]: ...
125
122
 
126
123
  child_elements: ClassVar[list[str]]
127
124
 
128
125
  def as_python_source(self) -> str:
129
126
  """Return the Python source code for this AST list."""
130
- mod = py_ast.Module(body=self.as_ast_list(), type_ignores=[], **DEFAULT_AST_ARGS_MODULE)
127
+ mod = py_ast.Module(body=self.as_ast_list(include_comments=True), type_ignores=[], **DEFAULT_AST_ARGS_MODULE)
131
128
  py_ast.fix_missing_locations(mod)
132
- return py_ast.unparse(mod)
129
+ return unparse_with_comments(mod)
133
130
 
134
131
 
135
132
  CodeGenAstType = CodeGenAst | CodeGenAstList
@@ -204,7 +201,7 @@ class Scope:
204
201
  return not self.is_name_reserved(name)
205
202
 
206
203
  while not _is_name_allowed(attempt):
207
- attempt = cleaned + str(count)
204
+ attempt = cleaned + "_" + str(count)
208
205
  count += 1
209
206
 
210
207
  return _add(attempt)
@@ -270,7 +267,7 @@ class _Annotation(Statement):
270
267
  self.name = name
271
268
  self.annotation = annotation
272
269
 
273
- def as_ast(self):
270
+ def as_ast(self, *, include_comments: bool = False):
274
271
  if not allowable_name(self.name):
275
272
  raise AssertionError(f"Expected {self.name} to be a valid Python identifier")
276
273
  return py_ast.AnnAssign(
@@ -293,7 +290,7 @@ class _Assignment(Statement):
293
290
  self.value = value
294
291
  self.type_hint = type_hint
295
292
 
296
- def as_ast(self):
293
+ def as_ast(self, *, include_comments: bool = False):
297
294
  if not allowable_name(self.name):
298
295
  raise AssertionError(f"Expected {self.name} to be a valid Python identifier")
299
296
  if self.type_hint is None:
@@ -320,32 +317,56 @@ class Block(CodeGenAstList):
320
317
 
321
318
  def __init__(self, scope: Scope, parent_block: Block | None = None):
322
319
  self.scope = scope
323
- # We all `Expression` here for things like MethodCall which
324
- # are bare expressions that are still useful for side effects
325
- self.statements: list[Block | Statement | Expression] = []
320
+ # We allow `Expression` here for things like MethodCall which
321
+ # are bare expressions that are still useful for side effects.
322
+ # `Comment` (str) entries are rendered as ``# text`` comments.
323
+ self.statements: list[Block | Statement | Expression | Comment] = []
326
324
  self.parent_block = parent_block
327
325
 
328
- def as_ast_list(self, allow_empty: bool = True) -> list[py_ast.stmt]:
326
+ def as_ast_list(self, allow_empty: bool = True, *, include_comments: bool = False) -> list[py_ast.stmt]:
329
327
  retval: list[py_ast.stmt] = []
330
328
  for s in self.statements:
329
+ if isinstance(s, str):
330
+ # Comment
331
+ if include_comments:
332
+ retval.append(CommentNode(s, **DEFAULT_AST_ARGS)) # type: ignore[reportArgumentType]
333
+ continue
331
334
  if isinstance(s, CodeGenAstList):
332
- retval.extend(s.as_ast_list(allow_empty=True))
335
+ retval.extend(s.as_ast_list(allow_empty=True, include_comments=include_comments))
336
+ elif isinstance(s, Statement):
337
+ ast_obj = s.as_ast(include_comments=include_comments)
338
+ assert isinstance(ast_obj, py_ast.stmt), (
339
+ "Statement object return {ast_obj} which is not a subclass of py_ast.stmt"
340
+ )
341
+ retval.append(ast_obj)
333
342
  else:
334
- if isinstance(s, Statement):
335
- ast_obj = s.as_ast()
336
- assert isinstance(ast_obj, py_ast.stmt), (
337
- "Statement object return {ast_obj} which is not a subclass of py_ast.stmt"
338
- )
339
- retval.append(ast_obj)
340
- else:
341
- # Things like bare function/method calls need to be wrapped
342
- # in `Expr` to match the way Python parses.
343
- retval.append(py_ast.Expr(s.as_ast(), **DEFAULT_AST_ARGS))
343
+ # Things like bare function/method calls need to be wrapped
344
+ # in `Expr` to match the way Python parses.
345
+ retval.append(py_ast.Expr(s.as_ast(include_comments=include_comments), **DEFAULT_AST_ARGS))
344
346
 
345
347
  if len(retval) == 0 and not allow_empty:
346
348
  return [py_ast.Pass(**DEFAULT_AST_ARGS)]
347
349
  return retval
348
350
 
351
+ def add_comment(self, text: str, *, wrap: int | None = None) -> None:
352
+ """Add a ``# text`` comment line at the current position in the block.
353
+
354
+ If *wrap* is given as an integer, long lines are wrapped at word
355
+ boundaries so that no comment line exceeds *wrap* characters
356
+ (including the ``# `` prefix).
357
+ """
358
+ if wrap is not None:
359
+ # Account for the "# " prefix (2 chars) that will be added during rendering.
360
+ effective_width = max(wrap - 2, 1)
361
+ lines = textwrap.wrap(text, width=effective_width)
362
+ if not lines:
363
+ # Empty or whitespace-only text – preserve as a single blank comment.
364
+ lines = [""]
365
+ for line in lines:
366
+ self.statements.append(line)
367
+ else:
368
+ self.statements.append(text)
369
+
349
370
  def add_statement(self, statement: Statement | Block | Expression) -> None:
350
371
  self.statements.append(statement)
351
372
  if isinstance(statement, Block):
@@ -509,6 +530,9 @@ class Block(CodeGenAstList):
509
530
  def create_return(self, value: Expression) -> None:
510
531
  self.add_statement(Return(value))
511
532
 
533
+ def create_assert(self, test: Expression, msg: Expression | None = None) -> None:
534
+ self.add_statement(Assert(test, msg))
535
+
512
536
  def create_if(self) -> If:
513
537
  """
514
538
  Create an If statement, add it to this block, and return it.
@@ -552,18 +576,16 @@ class Module(Block, CodeGenAst):
552
576
  for name in dir(builtins):
553
577
  scope.reserve_name(name, is_builtin=True)
554
578
  Block.__init__(self, scope)
555
- self.file_comments: list[str] = []
556
579
 
557
- def as_ast(self) -> py_ast.Module:
558
- return py_ast.Module(body=self.as_ast_list(), type_ignores=[], **DEFAULT_AST_ARGS_MODULE)
580
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.Module:
581
+ return py_ast.Module(
582
+ body=self.as_ast_list(include_comments=include_comments), type_ignores=[], **DEFAULT_AST_ARGS_MODULE
583
+ )
559
584
 
560
585
  def as_python_source(self) -> str:
561
- main = super().as_python_source()
562
- file_comments = "".join(f"# {comment}\n" for comment in self.file_comments)
563
- return file_comments + main
564
-
565
- def add_file_comment(self, comment: str) -> None:
566
- self.file_comments.append(comment)
586
+ mod = self.as_ast(include_comments=True)
587
+ py_ast.fix_missing_locations(mod)
588
+ return unparse_with_comments(mod)
567
589
 
568
590
 
569
591
  class ArgKind(enum.Enum):
@@ -677,7 +699,7 @@ class Function(Scope, Statement):
677
699
  self.reserve_name(arg.name, function_arg=True)
678
700
  self._args = combined
679
701
 
680
- def as_ast(self) -> py_ast.stmt:
702
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.stmt:
681
703
  if not allowable_name(self.func_name):
682
704
  raise AssertionError(f"Expected '{self.func_name}' to be a valid Python identifier")
683
705
  for arg in self._args:
@@ -716,7 +738,7 @@ class Function(Scope, Statement):
716
738
  defaults=defaults,
717
739
  **DEFAULT_AST_ARGS_ARGUMENTS,
718
740
  ),
719
- body=self.body.as_ast_list(allow_empty=False),
741
+ body=self.body.as_ast_list(allow_empty=False, include_comments=include_comments),
720
742
  decorator_list=[d.as_ast() for d in self.decorators],
721
743
  type_params=[],
722
744
  returns=self.return_type.as_ast() if self.return_type is not None else None,
@@ -743,14 +765,14 @@ class Class(Scope, Statement):
743
765
  self.bases: list[Expression] = list(bases) if bases else []
744
766
  self.decorators: list[Expression] = list(decorators) if decorators else []
745
767
 
746
- def as_ast(self) -> py_ast.stmt:
768
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.stmt:
747
769
  if not allowable_name(self.class_name):
748
770
  raise AssertionError(f"Expected '{self.class_name}' to be a valid Python identifier")
749
771
  return py_ast.ClassDef(
750
772
  name=self.class_name,
751
773
  bases=[b.as_ast() for b in self.bases],
752
774
  keywords=[],
753
- body=self.body.as_ast_list(allow_empty=False),
775
+ body=self.body.as_ast_list(allow_empty=False, include_comments=include_comments),
754
776
  decorator_list=[d.as_ast() for d in self.decorators],
755
777
  type_params=[],
756
778
  **DEFAULT_AST_ARGS,
@@ -763,13 +785,34 @@ class Return(Statement):
763
785
  def __init__(self, value: Expression):
764
786
  self.value = value
765
787
 
766
- def as_ast(self):
788
+ def as_ast(self, *, include_comments: bool = False):
767
789
  return py_ast.Return(self.value.as_ast(), **DEFAULT_AST_ARGS)
768
790
 
769
791
  def __repr__(self):
770
792
  return f"Return({repr(self.value)}"
771
793
 
772
794
 
795
+ class Assert(Statement):
796
+ """An ``assert`` statement with an optional message."""
797
+
798
+ child_elements = ["test", "msg"]
799
+
800
+ def __init__(self, test: Expression, msg: Expression | None = None):
801
+ self.test = test
802
+ self.msg = msg
803
+
804
+ def as_ast(self, *, include_comments: bool = False):
805
+ msg_ast = self.msg.as_ast() if self.msg is not None else None
806
+ return py_ast.Assert(
807
+ test=self.test.as_ast(),
808
+ msg=msg_ast,
809
+ **DEFAULT_AST_ARGS,
810
+ )
811
+
812
+ def __repr__(self):
813
+ return f"Assert({repr(self.test)}, {repr(self.msg)})"
814
+
815
+
773
816
  class If(Statement):
774
817
  child_elements = ["if_blocks", "conditions", "else_block"]
775
818
 
@@ -801,7 +844,7 @@ class If(Statement):
801
844
  return self.else_block
802
845
  return self
803
846
 
804
- def as_ast(self) -> py_ast.If:
847
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.If:
805
848
  if len(self.if_blocks) == 0:
806
849
  raise AssertionError("Should have called `finalize` on If")
807
850
  if_ast = empty_If()
@@ -809,7 +852,7 @@ class If(Statement):
809
852
  previous_if = None
810
853
  for condition, if_block in zip(self.conditions, self.if_blocks):
811
854
  current_if.test = condition.as_ast()
812
- current_if.body = if_block.as_ast_list()
855
+ current_if.body = if_block.as_ast_list(include_comments=include_comments)
813
856
  if previous_if is not None:
814
857
  previous_if.orelse.append(current_if)
815
858
 
@@ -818,7 +861,7 @@ class If(Statement):
818
861
 
819
862
  if self.else_block.statements:
820
863
  assert previous_if is not None
821
- previous_if.orelse = self.else_block.as_ast_list()
864
+ previous_if.orelse = self.else_block.as_ast_list(include_comments=include_comments)
822
865
 
823
866
  return if_ast
824
867
 
@@ -840,7 +883,7 @@ class With(Statement):
840
883
  self._parent_block = parent_block
841
884
  self.body = Block(parent_scope, parent_block=parent_block)
842
885
 
843
- def as_ast(self) -> py_ast.With:
886
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.With:
844
887
  optional_vars = None
845
888
  if self.target is not None:
846
889
  optional_vars = py_ast.Name(id=self.target.name, ctx=py_ast.Store(), **DEFAULT_AST_ARGS)
@@ -852,7 +895,7 @@ class With(Statement):
852
895
  optional_vars=optional_vars,
853
896
  )
854
897
  ],
855
- body=self.body.as_ast_list(allow_empty=False),
898
+ body=self.body.as_ast_list(allow_empty=False, include_comments=include_comments),
856
899
  **DEFAULT_AST_ARGS,
857
900
  )
858
901
 
@@ -866,9 +909,9 @@ class Try(Statement):
866
909
  self.except_block = Block(parent_scope)
867
910
  self.else_block = Block(parent_scope)
868
911
 
869
- def as_ast(self) -> py_ast.Try:
912
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.Try:
870
913
  return py_ast.Try(
871
- body=self.try_block.as_ast_list(allow_empty=False),
914
+ body=self.try_block.as_ast_list(allow_empty=False, include_comments=include_comments),
872
915
  handlers=[
873
916
  py_ast.ExceptHandler(
874
917
  type=(
@@ -881,11 +924,11 @@ class Try(Statement):
881
924
  )
882
925
  ),
883
926
  name=None,
884
- body=self.except_block.as_ast_list(allow_empty=False),
927
+ body=self.except_block.as_ast_list(allow_empty=False, include_comments=include_comments),
885
928
  **DEFAULT_AST_ARGS,
886
929
  )
887
930
  ],
888
- orelse=self.else_block.as_ast_list(allow_empty=True),
931
+ orelse=self.else_block.as_ast_list(allow_empty=True, include_comments=include_comments),
889
932
  finalbody=[],
890
933
  **DEFAULT_AST_ARGS,
891
934
  )
@@ -916,7 +959,7 @@ class Import(Statement):
916
959
  self.module = module
917
960
  self.as_ = as_
918
961
 
919
- def as_ast(self) -> py_ast.AST:
962
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.AST:
920
963
  if self.as_ is None:
921
964
  # No alias needed:
922
965
  return py_ast.Import(names=[py_ast.alias(name=self.module)], **DEFAULT_AST_ARGS)
@@ -941,7 +984,7 @@ class ImportFrom(Statement):
941
984
  self.import_ = import_
942
985
  self.as_ = as_
943
986
 
944
- def as_ast(self) -> py_ast.AST:
987
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.AST:
945
988
  if self.as_ is None:
946
989
  # No alias needed:
947
990
  return py_ast.ImportFrom(
@@ -961,7 +1004,7 @@ class ImportFrom(Statement):
961
1004
 
962
1005
  class Expression(CodeGenAst):
963
1006
  @abstractmethod
964
- def as_ast(self) -> py_ast.expr: ...
1007
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.expr: ...
965
1008
 
966
1009
  # Some utilities for easy chaining:
967
1010
 
@@ -983,6 +1026,9 @@ class Expression(CodeGenAst):
983
1026
  ) -> Call:
984
1027
  return self.attr(attribute).call(args, kwargs)
985
1028
 
1029
+ def subscript(self, slice: Expression, /) -> Subscript:
1030
+ return Subscript(self, slice)
1031
+
986
1032
  # Arithmetic operators
987
1033
 
988
1034
  def add(self, other: Expression, /) -> Add:
@@ -1054,12 +1100,10 @@ class Expression(CodeGenAst):
1054
1100
  class String(Expression):
1055
1101
  child_elements = []
1056
1102
 
1057
- type = str
1058
-
1059
1103
  def __init__(self, string_value: str):
1060
1104
  self.string_value = string_value
1061
1105
 
1062
- def as_ast(self) -> py_ast.expr:
1106
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.expr:
1063
1107
  return py_ast.Constant(
1064
1108
  self.string_value,
1065
1109
  kind=None, # 3.8, indicates no prefix, needed only for tests
@@ -1076,12 +1120,10 @@ class String(Expression):
1076
1120
  class Bool(Expression):
1077
1121
  child_elements = []
1078
1122
 
1079
- type = bool
1080
-
1081
1123
  def __init__(self, value: bool):
1082
1124
  self.value = value
1083
1125
 
1084
- def as_ast(self) -> py_ast.expr:
1126
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.expr:
1085
1127
  return py_ast.Constant(self.value, **DEFAULT_AST_ARGS)
1086
1128
 
1087
1129
  def __repr__(self):
@@ -1091,12 +1133,10 @@ class Bool(Expression):
1091
1133
  class Bytes(Expression):
1092
1134
  child_elements = []
1093
1135
 
1094
- type = bytes
1095
-
1096
1136
  def __init__(self, value: bytes):
1097
1137
  self.value = value
1098
1138
 
1099
- def as_ast(self) -> py_ast.expr:
1139
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.expr:
1100
1140
  return py_ast.Constant(self.value, **DEFAULT_AST_ARGS)
1101
1141
 
1102
1142
  def __repr__(self):
@@ -1109,7 +1149,7 @@ class Number(Expression):
1109
1149
  def __init__(self, number: int | float):
1110
1150
  self.number = number
1111
1151
 
1112
- def as_ast(self) -> py_ast.expr:
1152
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.expr:
1113
1153
  return py_ast.Constant(self.number, **DEFAULT_AST_ARGS)
1114
1154
 
1115
1155
  def __repr__(self):
@@ -1122,7 +1162,7 @@ class List(Expression):
1122
1162
  def __init__(self, items: list[Expression]):
1123
1163
  self.items = items
1124
1164
 
1125
- def as_ast(self) -> py_ast.expr:
1165
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.expr:
1126
1166
  return py_ast.List(elts=[i.as_ast() for i in self.items], ctx=py_ast.Load(), **DEFAULT_AST_ARGS)
1127
1167
 
1128
1168
 
@@ -1132,7 +1172,7 @@ class Tuple(Expression):
1132
1172
  def __init__(self, items: Sequence[Expression]):
1133
1173
  self.items = items
1134
1174
 
1135
- def as_ast(self) -> py_ast.expr:
1175
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.expr:
1136
1176
  return py_ast.Tuple(elts=[i.as_ast() for i in self.items], ctx=py_ast.Load(), **DEFAULT_AST_ARGS)
1137
1177
 
1138
1178
 
@@ -1142,7 +1182,7 @@ class Set(Expression):
1142
1182
  def __init__(self, items: Sequence[Expression]):
1143
1183
  self.items = items
1144
1184
 
1145
- def as_ast(self) -> py_ast.expr:
1185
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.expr:
1146
1186
  if len(self.items) == 0:
1147
1187
  # {} is a dict literal in Python, so empty sets must use set([])
1148
1188
  return py_ast.Call(
@@ -1160,7 +1200,7 @@ class Dict(Expression):
1160
1200
  def __init__(self, pairs: Sequence[tuple[Expression, Expression]]):
1161
1201
  self.pairs = pairs
1162
1202
 
1163
- def as_ast(self) -> py_ast.expr:
1203
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.expr:
1164
1204
  return py_ast.Dict(
1165
1205
  keys=[k.as_ast() for k, _ in self.pairs],
1166
1206
  values=[v.as_ast() for _, v in self.pairs],
@@ -1171,8 +1211,6 @@ class Dict(Expression):
1171
1211
  class StringJoinBase(Expression):
1172
1212
  child_elements = ["parts"]
1173
1213
 
1174
- type = str
1175
-
1176
1214
  def __init__(self, parts: Sequence[Expression]):
1177
1215
  self.parts = parts
1178
1216
 
@@ -1202,7 +1240,7 @@ class StringJoinBase(Expression):
1202
1240
 
1203
1241
 
1204
1242
  class FStringJoin(StringJoinBase):
1205
- def as_ast(self) -> py_ast.expr:
1243
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.expr:
1206
1244
  # f-strings
1207
1245
  values: list[py_ast.expr] = []
1208
1246
  for part in self.parts:
@@ -1221,7 +1259,7 @@ class FStringJoin(StringJoinBase):
1221
1259
 
1222
1260
 
1223
1261
  class ConcatJoin(StringJoinBase):
1224
- def as_ast(self) -> py_ast.expr:
1262
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.expr:
1225
1263
  # Concatenate with +
1226
1264
  left = self.parts[0].as_ast()
1227
1265
  for part in self.parts[1:]:
@@ -1249,7 +1287,7 @@ class Name(Expression):
1249
1287
  raise AssertionError(f"Cannot refer to undefined name '{name}'")
1250
1288
  self.name = name
1251
1289
 
1252
- def as_ast(self) -> py_ast.expr:
1290
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.expr:
1253
1291
  if not allowable_name(self.name, allow_builtin=True):
1254
1292
  raise AssertionError(f"Expected {self.name} to be a valid Python identifier")
1255
1293
  return py_ast.Name(id=self.name, ctx=py_ast.Load(), **DEFAULT_AST_ARGS)
@@ -1270,7 +1308,7 @@ class Attr(Expression):
1270
1308
  raise AssertionError(f"Expected {attribute} to be a valid Python identifier")
1271
1309
  self.attribute = attribute
1272
1310
 
1273
- def as_ast(self) -> py_ast.expr:
1311
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.expr:
1274
1312
  return py_ast.Attribute(value=self.value.as_ast(), attr=self.attribute, **DEFAULT_AST_ARGS)
1275
1313
 
1276
1314
 
@@ -1280,7 +1318,7 @@ class Starred(Expression):
1280
1318
  def __init__(self, value: Expression):
1281
1319
  self.value = value
1282
1320
 
1283
- def as_ast(self) -> py_ast.expr:
1321
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.expr:
1284
1322
  return py_ast.Starred(value=self.value.as_ast(), ctx=py_ast.Load(), **DEFAULT_AST_ARGS)
1285
1323
 
1286
1324
  def __repr__(self):
@@ -1314,7 +1352,7 @@ class Call(Expression):
1314
1352
  self.args = list(args)
1315
1353
  self.kwargs = kwargs
1316
1354
 
1317
- def as_ast(self) -> py_ast.expr:
1355
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.expr:
1318
1356
 
1319
1357
  for name in self.kwargs.keys():
1320
1358
  if not allowable_keyword_arg_name(name):
@@ -1373,17 +1411,17 @@ def method_call(
1373
1411
  return obj.attr(method_name).call(args=args, kwargs=kwargs)
1374
1412
 
1375
1413
 
1376
- class DictLookup(Expression):
1377
- child_elements = ["lookup_obj", "lookup_arg"]
1414
+ class Subscript(Expression):
1415
+ child_elements = ["value", "slice"]
1378
1416
 
1379
- def __init__(self, lookup_obj: Expression, lookup_arg: Expression):
1380
- self.lookup_obj = lookup_obj
1381
- self.lookup_arg = lookup_arg
1417
+ def __init__(self, value: Expression, slice: Expression):
1418
+ self.value = value
1419
+ self.slice = slice
1382
1420
 
1383
- def as_ast(self) -> py_ast.expr:
1421
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.expr:
1384
1422
  return py_ast.Subscript(
1385
- value=self.lookup_obj.as_ast(),
1386
- slice=py_ast.subscript_slice_object(self.lookup_arg.as_ast()),
1423
+ value=self.value.as_ast(),
1424
+ slice=py_ast.subscript_slice_object(self.slice.as_ast()),
1387
1425
  ctx=py_ast.Load(),
1388
1426
  **DEFAULT_AST_ARGS,
1389
1427
  )
@@ -1393,9 +1431,7 @@ create_class_instance = function_call
1393
1431
 
1394
1432
 
1395
1433
  class NoneExpr(Expression):
1396
- type = type(None)
1397
-
1398
- def as_ast(self) -> py_ast.expr:
1434
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.expr:
1399
1435
  return py_ast.Constant(value=None, **DEFAULT_AST_ARGS)
1400
1436
 
1401
1437
 
@@ -1412,7 +1448,7 @@ class ArithOp(BinaryOperator, ABC):
1412
1448
 
1413
1449
  op: ClassVar[type[py_ast.operator]]
1414
1450
 
1415
- def as_ast(self) -> py_ast.expr:
1451
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.expr:
1416
1452
  return py_ast.BinOp(
1417
1453
  left=self.left.as_ast(),
1418
1454
  op=self.op(**DEFAULT_AST_ARGS_ADD),
@@ -1456,10 +1492,9 @@ class MatMul(ArithOp):
1456
1492
  class CompareOp(BinaryOperator, ABC):
1457
1493
  """Comparison operator (ast.Compare)."""
1458
1494
 
1459
- type = bool
1460
1495
  op: ClassVar[type[py_ast.cmpop]]
1461
1496
 
1462
- def as_ast(self) -> py_ast.expr:
1497
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.expr:
1463
1498
  return py_ast.Compare(
1464
1499
  left=self.left.as_ast(),
1465
1500
  comparators=[self.right.as_ast()],
@@ -1501,10 +1536,9 @@ class NotIn(CompareOp):
1501
1536
 
1502
1537
 
1503
1538
  class BoolOp(BinaryOperator, ABC):
1504
- type = bool
1505
1539
  op: ClassVar[type[py_ast.boolop]]
1506
1540
 
1507
- def as_ast(self) -> py_ast.expr:
1541
+ def as_ast(self, *, include_comments: bool = False) -> py_ast.expr:
1508
1542
  return py_ast.BoolOp(
1509
1543
  op=self.op(),
1510
1544
  values=[self.left.as_ast(), self.right.as_ast()],
@@ -1631,7 +1665,7 @@ def auto(value: PythonObj) -> Expression:
1631
1665
  elif isinstance(value, (int, float)):
1632
1666
  return Number(value)
1633
1667
  elif value is None:
1634
- return NoneExpr()
1668
+ return constants.None_
1635
1669
  elif isinstance(value, list):
1636
1670
  return List([auto(item) for item in value])
1637
1671
  elif isinstance(value, tuple):
@@ -1648,6 +1682,6 @@ class constants:
1648
1682
  Useful pre-made Expression constants
1649
1683
  """
1650
1684
 
1651
- None_: NoneExpr = auto(None)
1685
+ None_: NoneExpr = NoneExpr()
1652
1686
  True_: Bool = auto(True)
1653
1687
  False_: Bool = auto(False)
File without changes
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: fluent-codegen
3
- Version: 0.1.0
3
+ Version: 0.2
4
4
  Summary: A Python library for generating Python code via AST construction.
5
5
  Author-email: Luke Plant <luke@lukeplant.me.uk>
6
6
  License-Expression: Apache-2.0
@@ -76,7 +76,7 @@ fluent method-chaining for expressions:
76
76
  func, _ = module.create_function("fizzbuzz", args=["n"])
77
77
 
78
78
  # 2. A Name reference to the "n" parameter (Function *is* a Scope)
79
- n = codegen.Name("n", func)
79
+ n = func.name("n")
80
80
 
81
81
  # 3. Build an if / elif / else chain
82
82
  if_stmt = func.body.create_if()
@@ -94,7 +94,7 @@ fluent method-chaining for expressions:
94
94
  branch.create_return(codegen.String("Buzz"))
95
95
 
96
96
  # else: return str(n)
97
- if_stmt.else_block.create_return(codegen.function_call("str", [n], {}, func))
97
+ if_stmt.else_block.create_return(module.scope.name("str").call([n]))
98
98
 
99
99
  # 4. Inspect the generated source
100
100
  print(module.as_python_source())
@@ -41,7 +41,7 @@ def test_reserve_name():
41
41
  name2 = scope.reserve_name("name")
42
42
  assert name1 == "name"
43
43
  assert name1 != name2
44
- assert name2 == "name2"
44
+ assert name2 == "name_2"
45
45
 
46
46
 
47
47
  def test_reserve_name_function_arg_disallowed():
@@ -81,7 +81,7 @@ def test_reserve_name_after_reserve_function_arg():
81
81
  scope = codegen.Scope()
82
82
  scope.reserve_function_arg_name("my_arg")
83
83
  name = scope.reserve_name("my_arg")
84
- assert name == "my_arg2"
84
+ assert name == "my_arg_2"
85
85
 
86
86
 
87
87
  def test_reserve_function_arg_after_reserve_name():
@@ -127,6 +127,47 @@ def test_function_return():
127
127
  )
128
128
 
129
129
 
130
+ def test_create_assert():
131
+ module = codegen.Module()
132
+ func = codegen.Function("myfunc", args=["x"], parent_scope=module.scope)
133
+ func.body.create_assert(func.name("x"))
134
+ assert_code_equal(
135
+ func,
136
+ """
137
+ def myfunc(x):
138
+ assert x
139
+ """,
140
+ )
141
+
142
+
143
+ def test_create_assert_with_message():
144
+ module = codegen.Module()
145
+ func = codegen.Function("myfunc", args=["x"], parent_scope=module.scope)
146
+ func.body.create_assert(func.name("x"), codegen.String("x must be truthy"))
147
+ assert_code_equal(
148
+ func,
149
+ """
150
+ def myfunc(x):
151
+ assert x, 'x must be truthy'
152
+ """,
153
+ )
154
+
155
+
156
+ def test_assert_class_directly():
157
+ assert_stmt = codegen.Assert(codegen.Number(1))
158
+ assert_code_equal(assert_stmt, "assert 1")
159
+
160
+
161
+ def test_assert_class_with_message_directly():
162
+ assert_stmt = codegen.Assert(codegen.Number(0), codegen.String("fail"))
163
+ assert_code_equal(assert_stmt, "assert 0, 'fail'")
164
+
165
+
166
+ def test_assert_repr():
167
+ assert_stmt = codegen.Assert(codegen.Number(1))
168
+ assert "Assert(" in repr(assert_stmt)
169
+
170
+
130
171
  def test_function_bad_name():
131
172
  module = codegen.Module()
132
173
  func = codegen.Function("my func", args=[], parent_scope=module.scope)
@@ -390,7 +431,7 @@ def test_create_annotation_duplicate_name_gets_renamed():
390
431
  module = codegen.Module()
391
432
  module.create_annotation("x", module.scope.name("int"))
392
433
  name2 = module.create_annotation("x", module.scope.name("str"))
393
- assert_code_equal(name2, "x2")
434
+ assert_code_equal(name2, "x_2")
394
435
 
395
436
 
396
437
  def test_create_annotation_bad_name():
@@ -470,7 +511,7 @@ def test_create_field_duplicate_name_gets_renamed():
470
511
  module = codegen.Module()
471
512
  module.create_field("x", module.scope.name("int"))
472
513
  name2 = module.create_field("x", module.scope.name("str"))
473
- assert_code_equal(name2, "x2")
514
+ assert_code_equal(name2, "x_2")
474
515
 
475
516
 
476
517
  # --- Function call tests ---
@@ -713,7 +754,7 @@ def test_if_scope():
713
754
  assert if_block.scope.is_name_in_use("myvalue")
714
755
 
715
756
  name_in_if_block = if_block.scope.reserve_name("myvalue")
716
- assert name_in_if_block == "myvalue2"
757
+ assert name_in_if_block == "myvalue_2"
717
758
 
718
759
 
719
760
  def test_block_create_if():
@@ -1043,10 +1084,10 @@ def test_string_join_collapse_strings():
1043
1084
  assert_code_equal(join1, "'hello there ' + tmp + ' how are you?'")
1044
1085
 
1045
1086
 
1046
- def test_dict_lookup():
1087
+ def test_subscript():
1047
1088
  scope = codegen.Scope()
1048
1089
  var = scope.create_name("tmp")
1049
- lookup = codegen.DictLookup(var, codegen.String("x"))
1090
+ lookup = var.subscript(codegen.String("x"))
1050
1091
  assert_code_equal(lookup, "tmp['x']")
1051
1092
 
1052
1093
 
@@ -1108,7 +1149,7 @@ def test_reserve_name_keyword_avoidance():
1108
1149
  scope = codegen.Scope()
1109
1150
  # 'class' is a keyword, so reserve_name should avoid it
1110
1151
  name = scope.reserve_name("class")
1111
- assert name == "class2" # skips 'class' because it's a keyword
1152
+ assert name == "class_2" # skips 'class' because it's a keyword
1112
1153
 
1113
1154
 
1114
1155
  def test_block_add_statement_bare_expression():
@@ -1172,7 +1213,6 @@ def test_dict_expression():
1172
1213
 
1173
1214
  def test_none_expr():
1174
1215
  n = codegen.NoneExpr()
1175
- assert n.type is type(None)
1176
1216
  assert_code_equal(n, "None")
1177
1217
 
1178
1218
 
@@ -1266,32 +1306,6 @@ def test_morph_into():
1266
1306
  assert s1.number == 42
1267
1307
 
1268
1308
 
1269
- def test_codegen_ast_not_implemented():
1270
- """Test that abstract methods raise NotImplementedError."""
1271
-
1272
- class DummyAst(codegen.CodeGenAst):
1273
- child_elements = []
1274
-
1275
- def as_ast(self):
1276
- raise NotImplementedError(f"{self.__class__!r}.as_ast()")
1277
-
1278
- d = DummyAst()
1279
- with pytest.raises(NotImplementedError):
1280
- d.as_ast()
1281
-
1282
-
1283
- def test_codegen_ast_list_not_implemented():
1284
- class DummyAstList(codegen.CodeGenAstList):
1285
- child_elements = []
1286
-
1287
- def as_ast_list(self, allow_empty=True):
1288
- raise NotImplementedError(f"{self.__class__!r}.as_ast_list()")
1289
-
1290
- d = DummyAstList()
1291
- with pytest.raises(NotImplementedError):
1292
- d.as_ast_list()
1293
-
1294
-
1295
1309
  def test_block_as_ast_list_with_codegen_ast_list():
1296
1310
  """Test that Block handles CodeGenAstList sub-statements."""
1297
1311
  module = codegen.Module()
@@ -1406,13 +1420,11 @@ def test_auto_unsupported_type():
1406
1420
 
1407
1421
  def test_bool_true():
1408
1422
  b = codegen.Bool(True)
1409
- assert b.type is bool
1410
1423
  assert_code_equal(b, "True")
1411
1424
 
1412
1425
 
1413
1426
  def test_bool_false():
1414
1427
  b = codegen.Bool(False)
1415
- assert b.type is bool
1416
1428
  assert_code_equal(b, "False")
1417
1429
 
1418
1430
 
@@ -1438,7 +1450,6 @@ def test_auto_bool_false():
1438
1450
 
1439
1451
  def test_bytes():
1440
1452
  b = codegen.Bytes(b"hello")
1441
- assert b.type is bytes
1442
1453
  assert_code_equal(b, "b'hello'")
1443
1454
 
1444
1455
 
@@ -1522,29 +1533,6 @@ def test_auto_frozenset():
1522
1533
  # The classes are tested indirectly by the utilty methods on `Expression`,
1523
1534
  # we don't really need separate tests.
1524
1535
 
1525
- # --- Comparison operator type ---
1526
-
1527
-
1528
- def test_comparison_ops_have_bool_type():
1529
- a, b = codegen.Number(1), codegen.Number(2)
1530
- for cls in (
1531
- codegen.Equals,
1532
- codegen.NotEquals,
1533
- codegen.Lt,
1534
- codegen.Gt,
1535
- codegen.LtE,
1536
- codegen.GtE,
1537
- codegen.In,
1538
- codegen.NotIn,
1539
- ):
1540
- assert cls(a, b).type is bool
1541
-
1542
-
1543
- def test_boolean_ops_have_bool_type():
1544
- a, b = codegen.Bool(True), codegen.Bool(False)
1545
- assert codegen.And(a, b).type is bool
1546
- assert codegen.Or(a, b).type is bool
1547
-
1548
1536
 
1549
1537
  # --- Expression utility method tests ---
1550
1538
 
@@ -2653,9 +2641,9 @@ def test_module_reserves_builtins_by_default():
2653
2641
  assert mod.scope.is_name_reserved("str")
2654
2642
 
2655
2643
 
2656
- def test_module_comment():
2644
+ def test_block_add_comment():
2657
2645
  mod = codegen.Module()
2658
- mod.add_file_comment("Hello")
2646
+ mod.add_comment("Hello")
2659
2647
  mod.create_import("foo")
2660
2648
  assert_code_equal(
2661
2649
  mod,
@@ -2664,3 +2652,160 @@ def test_module_comment():
2664
2652
  import foo
2665
2653
  """,
2666
2654
  )
2655
+
2656
+
2657
+ def test_comment_in_function_body():
2658
+ mod = codegen.Module()
2659
+ func, _ = mod.create_function("my_func", args=["x"])
2660
+ func.body.add_comment("Process the value")
2661
+ x = codegen.Name("x", func)
2662
+ func.body.create_return(x)
2663
+ assert_code_equal(
2664
+ mod,
2665
+ """
2666
+ def my_func(x):
2667
+ # Process the value
2668
+ return x
2669
+ """,
2670
+ )
2671
+
2672
+
2673
+ def test_comment_in_class_body():
2674
+ mod = codegen.Module()
2675
+ cls, _ = mod.create_class("MyClass")
2676
+ cls.body.add_comment("Class fields")
2677
+ cls.body.create_field("name", codegen.Name("str", mod.scope))
2678
+ assert_code_equal(
2679
+ mod,
2680
+ """
2681
+ class MyClass:
2682
+ # Class fields
2683
+ name: str
2684
+ """,
2685
+ )
2686
+
2687
+
2688
+ def test_multiple_comments():
2689
+ mod = codegen.Module()
2690
+ mod.add_comment("First comment")
2691
+ mod.add_comment("Second comment")
2692
+ mod.create_import("foo")
2693
+ assert_code_equal(
2694
+ mod,
2695
+ """
2696
+ # First comment
2697
+ # Second comment
2698
+ import foo
2699
+ """,
2700
+ )
2701
+
2702
+
2703
+ def test_comment_compile_constraint():
2704
+ """Comments must not affect compile() — as_ast() must return a compilable AST."""
2705
+ mod = codegen.Module()
2706
+ mod.add_comment("This is a comment")
2707
+ func, _ = mod.create_function("greet", args=["name"])
2708
+ func.body.add_comment("Return greeting")
2709
+ func.body.create_return(codegen.String("hello"))
2710
+
2711
+ # as_ast() must be compilable
2712
+ ast_tree = mod.as_ast()
2713
+ code = compile(ast_tree, "<test>", "exec")
2714
+ ns: dict[str, object] = {}
2715
+ exec(code, ns)
2716
+ assert ns["greet"]("world") == "hello" # type: ignore
2717
+
2718
+
2719
+ def test_comment_source_constraint():
2720
+ """as_python_source() must include the comments."""
2721
+ mod = codegen.Module()
2722
+ mod.add_comment("File header")
2723
+ mod.create_import("os")
2724
+ source = mod.as_python_source()
2725
+ assert "# File header" in source
2726
+ assert "import os" in source
2727
+
2728
+
2729
+ def test_comment_interleaved_with_statements():
2730
+ mod = codegen.Module()
2731
+ mod.create_import("os")
2732
+ mod.add_comment("Now import sys")
2733
+ mod.create_import("sys")
2734
+ assert_code_equal(
2735
+ mod,
2736
+ """
2737
+ import os
2738
+ # Now import sys
2739
+ import sys
2740
+ """,
2741
+ )
2742
+
2743
+
2744
+ def test_comment_only_block():
2745
+ """A block with only comments should compile (comments stripped = empty body)."""
2746
+ mod = codegen.Module()
2747
+ mod.add_comment("Just a comment")
2748
+ ast_tree = mod.as_ast()
2749
+ code = compile(ast_tree, "<test>", "exec")
2750
+ ns: dict[str, object] = {}
2751
+ exec(code, ns)
2752
+ # Should execute without error
2753
+
2754
+
2755
+ def test_add_comment_wrap_basic():
2756
+ mod = codegen.Module()
2757
+ mod.add_comment("This is a long comment that should be wrapped into multiple lines", wrap=40)
2758
+ mod.create_import("foo")
2759
+ assert_code_equal(
2760
+ mod,
2761
+ """
2762
+ # This is a long comment that should be
2763
+ # wrapped into multiple lines
2764
+ import foo
2765
+ """,
2766
+ )
2767
+
2768
+
2769
+ def test_add_comment_wrap_short_text_unchanged():
2770
+ """Text shorter than wrap width should remain a single comment."""
2771
+ mod = codegen.Module()
2772
+ mod.add_comment("Short", wrap=80)
2773
+ assert_code_equal(
2774
+ mod,
2775
+ """
2776
+ # Short
2777
+ """,
2778
+ )
2779
+
2780
+
2781
+ def test_add_comment_wrap_none_default():
2782
+ """Without wrap, long text stays on one line."""
2783
+ mod = codegen.Module()
2784
+ long_text = "word " * 30
2785
+ mod.add_comment(long_text.strip())
2786
+ source = mod.as_python_source()
2787
+ comment_lines = [line for line in source.splitlines() if line.startswith("#")]
2788
+ assert len(comment_lines) == 1
2789
+
2790
+
2791
+ def test_add_comment_wrap_in_function_body():
2792
+ mod = codegen.Module()
2793
+ func, _ = mod.create_function("f", args=[])
2794
+ func.body.add_comment("This is a very detailed explanation of what the function does", wrap=40)
2795
+ func.body.create_return(codegen.Number(1))
2796
+ assert_code_equal(
2797
+ mod,
2798
+ """
2799
+ def f():
2800
+ # This is a very detailed explanation of
2801
+ # what the function does
2802
+ return 1
2803
+ """,
2804
+ )
2805
+
2806
+
2807
+ def test_add_comment_wrap_empty_string():
2808
+ mod = codegen.Module()
2809
+ mod.add_comment("", wrap=40)
2810
+ source = mod.as_python_source()
2811
+ assert "#" in source
@@ -12,7 +12,7 @@ def _build_fizzbuzz_module() -> codegen.Module:
12
12
  module = codegen.Module()
13
13
  func, _ = module.create_function("fizzbuzz", args=["n"])
14
14
 
15
- n = codegen.Name("n", func)
15
+ n = func.name("n")
16
16
 
17
17
  if_stmt = func.body.create_if()
18
18
 
@@ -29,7 +29,7 @@ def _build_fizzbuzz_module() -> codegen.Module:
29
29
  branch_5.create_return(codegen.String("Buzz"))
30
30
 
31
31
  # else: return str(n)
32
- if_stmt.else_block.create_return(codegen.function_call("str", [n], {}, func))
32
+ if_stmt.else_block.create_return(module.scope.name("str").call([n]))
33
33
 
34
34
  return module
35
35
 
@@ -1 +0,0 @@
1
- __version__ = "0.1.0"
File without changes
File without changes