QuEraComputing#957 Improve docstrings

src/bloqade/builder/parse/builder.py

@@ -1,3 +1,14 @@
+"""
+Module for parsing quantum computing program components into intermediate representation (IR) using the bloqade library.
+
+This module provides a Parser class for parsing various components of a quantum computing program, including atom arrangements,
+pulse sequences, analog circuits, and routines. It also defines utility functions for reading addresses, waveforms, drives,
+sequences, registers, and pragmas from a builder stream.
+
+Classes:
+    Parser: A class for parsing quantum computing program components into intermediate representation (IR).
+"""
+
 from bloqade.builder.base import Builder
 from bloqade.builder.coupling import LevelCoupling, Rydberg, Hyperfine
 from bloqade.builder.sequence_builder import SequenceBuilder
@@ -16,8 +27,9 @@
     from bloqade.ir.routine.base import Routine
     from bloqade.ir.analog_circuit import AnalogCircuit
 
-
 class Parser:
+    """A class for parsing quantum computing program components into intermediate representation (IR)."""
+
     stream: Optional["BuilderStream"] = None
     vector_node_names: Set[str] = set()
     sequence: ir.Sequence = ir.Sequence.create()
@@ -27,6 +39,7 @@ class Parser:
     order: Tuple[str, ...] = ()
 
     def reset(self, builder: Builder):
+        """Reset the parser's state."""
         self.stream = BuilderStream.create(builder)
         self.vector_node_names = set()
         self.sequence = ir.Sequence.create()
@@ -36,6 +49,15 @@ def reset(self, builder: Builder):
         self.order = ()
 
     def read_address(self, stream) -> Tuple[LevelCoupling, Field, BuilderNode]:
+        """
+        Read an address from the builder stream.
+
+        Args:
+            stream: The builder stream.
+
+        Returns:
+            Tuple[LevelCoupling, Field, BuilderNode]: A tuple containing the level coupling, field, and spatial modulation.
+        """
         spatial = stream.read_next([Location, Uniform, Scale])
         curr = spatial
 
@@ -62,6 +84,15 @@ def read_address(self, stream) -> Tuple[LevelCoupling, Field, BuilderNode]:
             return (None, None, spatial)
 
     def read_waveform(self, head: BuilderNode) -> Tuple[ir.Waveform, BuilderNode]:
+        """
+        Read a waveform from the builder stream.
+
+        Args:
+            head (BuilderNode): The head of the builder stream.
+
+        Returns:
+            Tuple[ir.Waveform, BuilderNode]: A tuple containing the waveform and the next builder node.
+        """
         curr = head
         waveform = None
         while curr is not None:
@@ -103,6 +134,15 @@ def read_waveform(self, head: BuilderNode) -> Tuple[ir.Waveform, BuilderNode]:
         return waveform, curr
 
     def read_drive(self, head) -> ir.Field:
+        """
+        Read a drive from the builder stream.
+
+        Args:
+            head: The head of the builder stream.
+
+        Returns:
+            ir.Field: The drive field.
+        """
         if head is None:
             return ir.Field({})
 
@@ -112,6 +152,12 @@ def read_drive(self, head) -> ir.Field:
         return ir.Field({sm: wf})
 
     def read_sequence(self) -> ir.Sequence:
+        """
+        Read a sequence from the builder stream.
+
+        Returns:
+            ir.Sequence: The parsed sequence.
+        """
         if isinstance(self.stream.curr.node, SequenceBuilder):
             # case with sequence builder object.
             self.sequence = self.stream.read().node._sequence
@@ -146,13 +192,20 @@ def read_sequence(self) -> ir.Sequence:
         return self.sequence
 
     def read_register(self) -> ir.AtomArrangement:
+        """
+        Read an atom arrangement register from the builder stream.
+
+        Returns:
+            ir.AtomArrangement: The parsed atom arrangement.
+        """
         # register is always head of the stream
         register_node = self.stream.read()
         self.register = register_node.node
 
         return self.register
 
     def read_pragmas(self) -> None:
+        """Read pragmas from the builder stream."""
         pragma_types = (
             Assign,
             BatchAssign,
@@ -199,17 +252,44 @@ def read_pragmas(self) -> None:
     def parse_register(
         self, builder: Builder
     ) -> Union[ir.AtomArrangement, ir.ParallelRegister]:
+        """
+        Parse an atom arrangement register from the builder.
+
+        Args:
+            builder (Builder): The builder instance.
+
+        Returns:
+            Union[ir.AtomArrangement, ir.ParallelRegister]: The parsed atom arrangement or parallel register.
+        """
         self.reset(builder)
         self.read_register()
         self.read_pragmas()
         return self.register
 
     def parse_sequence(self, builder: Builder) -> ir.Sequence:
+        """
+        Parse a sequence from the builder.
+
+        Args:
+            builder (Builder): The builder instance.
+
+        Returns:
+            ir.Sequence: The parsed sequence.
+        """
         self.reset(builder)
         self.read_sequence()
         return self.sequence
 
     def parse_circuit(self, builder: Builder) -> "AnalogCircuit":
+        """
+        Parse an analog circuit from the builder.
+
+        Args:
+            builder (Builder): The builder instance.
+
+        Returns:
+            AnalogCircuit: The parsed analog circuit.
+        """
         from bloqade.ir.analog_circuit import AnalogCircuit
 
         self.reset(builder)
@@ -221,6 +301,15 @@ def parse_circuit(self, builder: Builder) -> "AnalogCircuit":
         return circuit
 
     def parse(self, builder: Builder) -> "Routine":
+        """
+        Parse a routine from the builder.
+
+        Args:
+            builder (Builder): The builder instance.
+
+        Returns:
+            Routine: The parsed routine.
+        """
         from bloqade.ir.analog_circuit import AnalogCircuit
         from bloqade.ir.routine.params import Params, ScalarArg, VectorArg
         from bloqade.ir.routine.base import Routine

src/bloqade/builder/parse/stream.py

@@ -1,17 +1,30 @@
+"""
+Module for managing a stream of builder nodes.
+
+This module provides classes to represent builder nodes and builder streams. A builder node is a single
+element in the stream, representing a step in a construction process. A builder stream is a sequence
+of builder nodes, allowing traversal and manipulation of the construction steps.
+
+Classes:
+    BuilderNode: Represents a single node in the builder stream.
+    BuilderStream: Represents a stream of builder nodes.
+"""
+
 from dataclasses import dataclass
 from typing import Optional, List, Type
 from bloqade.builder.base import Builder
 
-
 @dataclass
 class BuilderNode:
+    """A node in the builder stream."""
+
     node: Builder
     next: Optional["BuilderNode"] = None
 
     def __repr__(self) -> str:
+        """Representation of the BuilderNode."""
         return repr(self.node)
 
-
 @dataclass
 class BuilderStream:
     """Represents a stream of builder nodes."""
@@ -20,9 +33,11 @@ class BuilderStream:
     curr: Optional[BuilderNode] = None
 
     def copy(self) -> "BuilderStream":
+        """Create a copy of the builder stream."""
         return BuilderStream(head=self.head, curr=self.curr)
 
     def read(self) -> Optional[BuilderNode]:
+        """Read the next builder node from the stream."""
         if self.curr is None:
             return None
 
@@ -31,6 +46,15 @@ def read(self) -> Optional[BuilderNode]:
         return node
 
     def read_next(self, builder_types: List[type[Builder]]) -> Optional[BuilderNode]:
+        """
+        Read the next builder node of specified types from the stream.
+
+        Args:
+            builder_types (List[type[Builder]]): List of builder types to read from the stream.
+
+        Returns:
+            Optional[BuilderNode]: The next builder node matching one of the specified types, or None if not found.
+        """
         node = self.read()
         while node is not None:
             if type(node.node) in builder_types:
@@ -41,12 +65,12 @@ def read_next(self, builder_types: List[type[Builder]]) -> Optional[BuilderNode]
     def eat(
         self, types: List[Type[Builder]], skips: Optional[List[Type[Builder]]] = None
     ) -> BuilderNode:
-        """Scan the stream until a node of type in `types` or `skips` is found.
+        """
+        Move the stream pointer until a node of specified types is found.
 
         Args:
-            types (List[Type[Builder]]): List of types to move the stream pointer to
-            skips (List[Type[Builder]] | None, optional): List of types to end the
-            stream scan
+            types (List[Type[Builder]]): List of types to move the stream pointer to.
+            skips (List[Type[Builder]] | None, optional): List of types to end the stream scan.
 
         Returns:
             BuilderNode: The beginning of the stream which matches a type in `types`.
@@ -62,16 +86,27 @@ def eat(
         return head
 
     def __iter__(self):
+        """Iterator method to iterate over the builder stream."""
         return self
 
     def __next__(self):
+        """Next method to get the next item in the builder stream."""
         node = self.read()
         if node is None:
             raise StopIteration
         return node
 
     @staticmethod
     def build_nodes(node: Builder) -> "BuilderNode":
+        """
+        Build BuilderNode instances from the provided Builder.
+
+        Args:
+            node (Builder): The root Builder instance.
+
+        Returns:
+            BuilderNode: The head of the linked list of BuilderNodes.
+        """
         curr = node
         node = None
         while curr is not None:
@@ -83,5 +118,14 @@ def build_nodes(node: Builder) -> "BuilderNode":
 
     @staticmethod
     def create(builder: Builder) -> "BuilderStream":
+        """
+        Create a BuilderStream instance from a Builder.
+
+        Args:
+            builder (Builder): The root Builder instance.
+
+        Returns:
+            BuilderStream: The created BuilderStream instance.
+        """
         head = BuilderStream.build_nodes(builder)
         return BuilderStream(head=head, curr=head)