Support pydantic v2 and get rid of plum warnings (#922)

* Support pydantic v2 the easiest way possible and get rid of plum warnings

* whoops forgot lint

* bump version

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bloqade"
-version = "0.15.7"
+version = "0.15.8"
 description = "Neutral atom software development kit"
 authors = [
     {name = "QuEra Computing Inc.", email = "[email protected]"},
@@ -18,7 +18,7 @@ classifiers = [
 dependencies = [
     "juliacall>=0.9.14",
     "numpy>=1.25.2",
-    "pydantic>=1.10.13",
+    "pydantic>=2.0",
     "scipy>=1.9.3",
     "pandas>=2.1.0",
     "bokeh>=3.2.2",

src/bloqade/compiler/analysis/common/scan_variables.py

@@ -1,4 +1,4 @@
-from pydantic.dataclasses import dataclass
+from pydantic.v1.dataclasses import dataclass
 import bloqade.ir.scalar as scalar
 import bloqade.ir.control.field as field
 

src/bloqade/compiler/analysis/python/waveform.py

@@ -2,7 +2,7 @@
 import bloqade.ir.control.waveform as waveform
 from beartype.typing import Dict, FrozenSet, Set
 from beartype import beartype
-from pydantic.dataclasses import dataclass
+from pydantic.v1.dataclasses import dataclass
 
 
 @dataclass(frozen=True)

src/bloqade/compiler/codegen/hardware/lattice.py

@@ -6,7 +6,7 @@
 from beartype.typing import Optional
 import numpy as np
 from decimal import Decimal
-from pydantic.dataclasses import dataclass
+from pydantic.v1.dataclasses import dataclass
 from beartype.typing import List, Tuple
 
 

src/bloqade/compiler/codegen/hardware/piecewise_constant.py

@@ -6,7 +6,7 @@
 
 from beartype.typing import List
 from beartype import beartype
-from pydantic.dataclasses import dataclass
+from pydantic.v1.dataclasses import dataclass
 from bisect import bisect_right, bisect_left
 from decimal import Decimal
 

src/bloqade/compiler/codegen/hardware/piecewise_linear.py

@@ -5,7 +5,7 @@
 
 from beartype.typing import List
 from beartype import beartype
-from pydantic.dataclasses import dataclass
+from pydantic.v1.dataclasses import dataclass
 from bisect import bisect_left, bisect_right
 from decimal import Decimal
 

src/bloqade/compiler/passes/hardware/components.py

@@ -2,7 +2,7 @@
 from bloqade.compiler.codegen.hardware.piecewise_linear import PiecewiseLinear
 from bloqade.compiler.codegen.hardware.piecewise_constant import PiecewiseConstant
 from beartype.typing import Optional, List
-from pydantic.dataclasses import dataclass
+from pydantic.v1.dataclasses import dataclass
 from decimal import Decimal
 
 

src/bloqade/ir/analog_circuit.py

@@ -4,7 +4,7 @@
 from bloqade.ir.location.location import AtomArrangement, ParallelRegister
 from bloqade.ir.tree_print import Printer
 from beartype.typing import Union
-from pydantic.dataclasses import dataclass
+from pydantic.v1.dataclasses import dataclass
 
 
 # NOTE: this is just a dummy type bundle geometry and sequence

src/bloqade/ir/control/field.py

@@ -4,7 +4,7 @@
 from bloqade.ir.control.waveform import Waveform
 from bloqade.ir.control.traits import HashTrait, CanonicalizeTrait
 from bloqade.visualization import get_field_figure
-from pydantic.dataclasses import dataclass
+from pydantic.v1.dataclasses import dataclass
 from beartype.typing import Dict, List, Optional
 from decimal import Decimal
 from bloqade.visualization import display_ir

src/bloqade/ir/control/pulse.py

@@ -11,7 +11,7 @@
     CanonicalizeTrait,
 )
 from beartype.typing import List
-from pydantic.dataclasses import dataclass
+from pydantic.v1.dataclasses import dataclass
 from bloqade.visualization import get_pulse_figure
 from bloqade.visualization import display_ir
 

src/bloqade/ir/control/sequence.py

@@ -10,7 +10,7 @@
 from bloqade.ir.scalar import Interval, Scalar, cast
 from bloqade.ir.tree_print import Printer
 
-from pydantic.dataclasses import dataclass
+from pydantic.v1.dataclasses import dataclass
 from beartype.typing import List, Dict
 from bloqade.visualization import get_ir_figure
 from bloqade.visualization import display_ir

src/bloqade/ir/control/waveform.py

@@ -19,7 +19,7 @@
 
 from bisect import bisect_left, bisect_right
 from decimal import Decimal
-from pydantic.dataclasses import dataclass
+from pydantic.v1.dataclasses import dataclass
 from beartype.typing import Any, Tuple, Union, List, Callable, Dict, Container
 from beartype import beartype
 from enum import Enum

src/bloqade/ir/location/bravais.py

@@ -1,5 +1,5 @@
 from functools import cached_property
-from pydantic.dataclasses import dataclass
+from pydantic.v1.dataclasses import dataclass
 from dataclasses import fields
 import numpy as np
 import itertools

src/bloqade/ir/location/location.py

@@ -3,8 +3,9 @@
 from bloqade.ir.scalar import Scalar, Literal, cast
 from bloqade.ir.tree_print import Printer
 
-from pydantic.dataclasses import dataclass
+from pydantic.v1.dataclasses import dataclass
 from beartype.typing import List, Tuple, Generator, Union, Optional
+from beartype.door import is_bearable
 from beartype import beartype
 from enum import Enum
 from numpy.typing import NDArray
@@ -14,7 +15,6 @@
 
 from beartype.vale import Is
 from typing import Annotated
-from plum import dispatch
 import plotext as pltxt
 import sys
 import numpy as np
@@ -234,8 +234,8 @@ def scale(self, scale: ScalarType):
 
         return ListOfLocations(location_list)
 
-    @dispatch
-    def _add_position(
+    ## enforce input is single tuple and filling is single boolean
+    def add_position_single_tupe(
         self, position: Tuple[ScalarType, ScalarType], filling: Optional[bool] = None
     ):
         if filling is None:
@@ -246,8 +246,8 @@ def _add_position(
 
         return ListOfLocations(location_list)
 
-    @dispatch
-    def _add_position(  # noqa: F811
+    ## enforce input is list of tuples and filling is list of booleans or none
+    def add_position_list_tuples(  # noqa: F811
         self,
         position: List[Tuple[ScalarType, ScalarType]],
         filling: Optional[List[bool]] = None,
@@ -268,8 +268,8 @@ def _add_position(  # noqa: F811
 
         return ListOfLocations(location_list)
 
-    @dispatch
-    def _add_position(  # noqa: F811
+    ## enforce input is numpy array and filling is numpy array or none
+    def add_position_ndarray(  # noqa: F811
         self, position: PositionArray, filling: Optional[BoolArray] = None
     ):
         return self.add_position(
@@ -333,7 +333,21 @@ def add_position(
             shows your geometry in your web browser
 
         """
-        return self._add_position(position, filling)
+
+        if is_bearable(position, PositionArray) and is_bearable(
+            filling, Optional[BoolArray]
+        ):
+            return self.add_position_ndarray(position, filling)
+        elif is_bearable(position, List[Tuple[ScalarType, ScalarType]]) and is_bearable(
+            filling, Optional[List[bool]]
+        ):
+            return self.add_position_list_tuples(position, filling)
+        elif is_bearable(position, Tuple[ScalarType, ScalarType]) and is_bearable(
+            filling, Optional[bool]
+        ):
+            return self.add_position_single_tupe(position, filling)
+        else:
+            raise TypeError("Invalid input types for add_position provided!")
 
     @beartype
     def apply_defect_count(

src/bloqade/ir/routine/base.py

@@ -6,8 +6,8 @@
 from bloqade.builder.parse.trait import Parse, Show
 from bloqade.ir.routine.params import Params
 
-from pydantic import ConfigDict
-from pydantic.dataclasses import dataclass
+from pydantic.v1 import ConfigDict
+from pydantic.v1.dataclasses import dataclass
 from typing import TYPE_CHECKING, Union, Optional
 
 if TYPE_CHECKING:

src/bloqade/ir/routine/bloqade.py

@@ -5,7 +5,7 @@
 from bloqade.task.batch import LocalBatch
 from beartype import beartype
 from beartype.typing import Optional, Tuple, Callable, Dict, Any, List, NamedTuple
-from pydantic.dataclasses import dataclass
+from pydantic.v1.dataclasses import dataclass
 import numpy as np
 
 from bloqade.emulate.codegen.hamiltonian import CompileCache, RydbergHamiltonianCodeGen

src/bloqade/ir/routine/braket.py

@@ -1,5 +1,5 @@
 from collections import OrderedDict
-from pydantic.dataclasses import dataclass
+from pydantic.v1.dataclasses import dataclass
 from beartype import beartype
 from beartype.typing import Optional, Tuple
 from bloqade.builder.typing import LiteralType

src/bloqade/ir/routine/params.py

@@ -1,5 +1,5 @@
 from functools import cached_property
-from pydantic.dataclasses import dataclass
+from pydantic.v1.dataclasses import dataclass
 from typing import Dict, List, Tuple, Union
 from decimal import Decimal
 

src/bloqade/ir/routine/quera.py

@@ -1,5 +1,5 @@
 from collections import OrderedDict
-from pydantic.dataclasses import dataclass
+from pydantic.v1.dataclasses import dataclass
 import json
 
 from bloqade.builder.typing import LiteralType

src/bloqade/ir/scalar.py

@@ -1,7 +1,7 @@
 from typing import Any, Optional
 import numpy as np
-from pydantic.dataclasses import dataclass
-from pydantic import ValidationError, validator
+from pydantic.v1.dataclasses import dataclass
+from pydantic.v1 import ValidationError, validator
 from .tree_print import Printer
 import re
 from decimal import Decimal

src/bloqade/ir/tree_print.py

@@ -1,6 +1,6 @@
 # Ported from the Julia language AbstractTrees.jl implementation: https://github.com/JuliaCollections/AbstractTrees.jl/blob/master/src/printing.jl
 
-from pydantic.dataclasses import dataclass
+from pydantic.v1.dataclasses import dataclass
 import sys
 import bloqade.ir as ir
 

src/bloqade/submission/base.py

@@ -2,7 +2,7 @@
 from bloqade.submission.ir.task_specification import QuEraTaskSpecification
 from bloqade.submission.ir.task_results import QuEraTaskResults, QuEraTaskStatusCode
 from typing import Union
-from pydantic import BaseModel, Extra
+from pydantic.v1 import BaseModel, Extra
 from bloqade.submission.capabilities import get_capabilities
 from bloqade.submission.ir.capabilities import QuEraCapabilities
 

src/bloqade/submission/braket.py

@@ -14,7 +14,7 @@
 from bloqade.submission.ir.task_specification import QuEraTaskSpecification
 from braket.aws import AwsDevice, AwsQuantumTask
 from beartype.typing import Optional
-from pydantic import PrivateAttr
+from pydantic.v1 import PrivateAttr
 import bloqade
 
 

src/bloqade/submission/ir/braket.py

@@ -21,7 +21,7 @@
     LocalField,
 )
 from typing import Tuple, Union, List
-from pydantic import BaseModel
+from pydantic.v1 import BaseModel
 from decimal import Decimal
 
 

src/bloqade/submission/ir/capabilities.py

@@ -1,4 +1,4 @@
-from pydantic import BaseModel
+from pydantic.v1 import BaseModel
 from decimal import Decimal
 from beartype.typing import Optional
 from beartype import beartype

src/bloqade/submission/ir/parallel.py

@@ -1,4 +1,4 @@
-from pydantic import BaseModel, validator, ValidationError
+from pydantic.v1 import BaseModel, validator, ValidationError
 
 from typing import List, Optional, Tuple
 from itertools import combinations

src/bloqade/submission/ir/task_results.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from pydantic import BaseModel, conlist, conint
+from pydantic.v1 import BaseModel, conlist, conint
 from typing import List, Tuple
 import numpy as np
 

src/bloqade/submission/ir/task_specification.py

@@ -1,4 +1,4 @@
-from pydantic import BaseModel
+from pydantic.v1 import BaseModel
 from typing import Optional, List, Tuple
 from decimal import Decimal
 from bloqade.submission.ir.capabilities import QuEraCapabilities

src/bloqade/submission/quera.py

@@ -1,4 +1,4 @@
-from pydantic import PrivateAttr
+from pydantic.v1 import PrivateAttr
 from bloqade.submission.base import SubmissionBackend, ValidationError
 from bloqade.submission.ir.task_specification import (
     QuEraTaskSpecification,

src/bloqade/task/base.py

@@ -11,7 +11,7 @@
 from numpy.typing import NDArray
 import pandas as pd
 import numpy as np
-from pydantic.dataclasses import dataclass
+from pydantic.v1.dataclasses import dataclass
 from bloqade.submission.ir.parallel import ParallelDecoder
 import datetime
 from bloqade.visualization import display_report

tests/test_builder.py

@@ -12,7 +12,6 @@
 # import bloqade.builder.backend as builder_backend
 import bloqade.ir.routine.quera as quera
 import bloqade.ir.routine.braket as braket
-from plum import NotFoundLookupError
 
 from bloqade.ir.control.waveform import to_waveform
 from bloqade.ir import rydberg, detuning, hyperfine, rabi
@@ -77,13 +76,13 @@ def test_add_position_dispatch():
     with pytest.raises(AssertionError):
         start.add_position(position_list, [True])
 
-    with pytest.raises(NotFoundLookupError):
+    with pytest.raises(TypeError):
         start.add_position(position_list, True)
 
-    with pytest.raises(NotFoundLookupError):
+    with pytest.raises(TypeError):
         start.add_position(position_list, np.array([True, True]))
 
-    with pytest.raises(NotFoundLookupError):
+    with pytest.raises(TypeError):
         start.add_position(position, [True, True])
 
 

tests/test_scalar.py

@@ -1,4 +1,4 @@
-from pydantic import ValidationError
+from pydantic.v1 import ValidationError
 from bloqade import cast, var
 import bloqade.ir.scalar as scalar
 import pytest