From 05a9f94f2cb99665523eb75d4506771b12ca8b15 Mon Sep 17 00:00:00 2001
From: John Long <johnzl@protonmail.ch>
Date: Tue, 2 Apr 2024 10:58:55 -0400
Subject: [PATCH] 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
---
 pdm.lock                                      |  2 +-
 pyproject.toml                                |  4 +--
 .../analysis/common/scan_variables.py         |  2 +-
 .../compiler/analysis/python/waveform.py      |  2 +-
 .../compiler/codegen/hardware/lattice.py      |  2 +-
 .../codegen/hardware/piecewise_constant.py    |  2 +-
 .../codegen/hardware/piecewise_linear.py      |  2 +-
 .../compiler/passes/hardware/components.py    |  2 +-
 src/bloqade/ir/analog_circuit.py              |  2 +-
 src/bloqade/ir/control/field.py               |  2 +-
 src/bloqade/ir/control/pulse.py               |  2 +-
 src/bloqade/ir/control/sequence.py            |  2 +-
 src/bloqade/ir/control/waveform.py            |  2 +-
 src/bloqade/ir/location/bravais.py            |  2 +-
 src/bloqade/ir/location/location.py           | 32 +++++++++++++------
 src/bloqade/ir/routine/base.py                |  4 +--
 src/bloqade/ir/routine/bloqade.py             |  2 +-
 src/bloqade/ir/routine/braket.py              |  2 +-
 src/bloqade/ir/routine/params.py              |  2 +-
 src/bloqade/ir/routine/quera.py               |  2 +-
 src/bloqade/ir/scalar.py                      |  4 +--
 src/bloqade/ir/tree_print.py                  |  2 +-
 src/bloqade/submission/base.py                |  2 +-
 src/bloqade/submission/braket.py              |  2 +-
 src/bloqade/submission/ir/braket.py           |  2 +-
 src/bloqade/submission/ir/capabilities.py     |  2 +-
 src/bloqade/submission/ir/parallel.py         |  2 +-
 src/bloqade/submission/ir/task_results.py     |  2 +-
 .../submission/ir/task_specification.py       |  2 +-
 src/bloqade/submission/quera.py               |  2 +-
 src/bloqade/task/base.py                      |  2 +-
 tests/test_builder.py                         |  7 ++--
 tests/test_scalar.py                          |  2 +-
 33 files changed, 60 insertions(+), 47 deletions(-)

diff --git a/pdm.lock b/pdm.lock
index 4d6906e20..785a27f4e 100644
--- a/pdm.lock
+++ b/pdm.lock
@@ -5,7 +5,7 @@
 groups = ["default", "dev", "doc"]
 strategy = ["cross_platform", "inherit_metadata"]
 lock_version = "4.4.1"
-content_hash = "sha256:553f474cf8eaff800f485e34ba0b41ddc4a5aa6a32f782467c9672f86a43ca00"
+content_hash = "sha256:e6814e52f0916471a3e66e84cc7db1df2eb0e7fbdc579c2782430a47669a19bb"
 
 [[package]]
 name = "amazon-braket-default-simulator"
diff --git a/pyproject.toml b/pyproject.toml
index 1f2eb0b3b..7ccc5cb50 100644
--- a/pyproject.toml
+++ b/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 = "info@quera.com"},
@@ -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",
diff --git a/src/bloqade/compiler/analysis/common/scan_variables.py b/src/bloqade/compiler/analysis/common/scan_variables.py
index 0149e552f..712461ef9 100644
--- a/src/bloqade/compiler/analysis/common/scan_variables.py
+++ b/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
 
diff --git a/src/bloqade/compiler/analysis/python/waveform.py b/src/bloqade/compiler/analysis/python/waveform.py
index b15c6a149..2fada2b80 100644
--- a/src/bloqade/compiler/analysis/python/waveform.py
+++ b/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)
diff --git a/src/bloqade/compiler/codegen/hardware/lattice.py b/src/bloqade/compiler/codegen/hardware/lattice.py
index f5e8c679f..3f7eab124 100644
--- a/src/bloqade/compiler/codegen/hardware/lattice.py
+++ b/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
 
 
diff --git a/src/bloqade/compiler/codegen/hardware/piecewise_constant.py b/src/bloqade/compiler/codegen/hardware/piecewise_constant.py
index 6af8774eb..09fae079c 100644
--- a/src/bloqade/compiler/codegen/hardware/piecewise_constant.py
+++ b/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
 
diff --git a/src/bloqade/compiler/codegen/hardware/piecewise_linear.py b/src/bloqade/compiler/codegen/hardware/piecewise_linear.py
index fd821bd13..dbb6a181f 100644
--- a/src/bloqade/compiler/codegen/hardware/piecewise_linear.py
+++ b/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
 
diff --git a/src/bloqade/compiler/passes/hardware/components.py b/src/bloqade/compiler/passes/hardware/components.py
index a1ccc0b28..92b74a2f4 100644
--- a/src/bloqade/compiler/passes/hardware/components.py
+++ b/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
 
 
diff --git a/src/bloqade/ir/analog_circuit.py b/src/bloqade/ir/analog_circuit.py
index 3461b0950..d10d9b658 100644
--- a/src/bloqade/ir/analog_circuit.py
+++ b/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
diff --git a/src/bloqade/ir/control/field.py b/src/bloqade/ir/control/field.py
index 3a24e4f8b..2a0b65ca5 100644
--- a/src/bloqade/ir/control/field.py
+++ b/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
diff --git a/src/bloqade/ir/control/pulse.py b/src/bloqade/ir/control/pulse.py
index 118346671..47e9fec23 100644
--- a/src/bloqade/ir/control/pulse.py
+++ b/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
 
diff --git a/src/bloqade/ir/control/sequence.py b/src/bloqade/ir/control/sequence.py
index 326c6293a..83d5ac119 100644
--- a/src/bloqade/ir/control/sequence.py
+++ b/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
diff --git a/src/bloqade/ir/control/waveform.py b/src/bloqade/ir/control/waveform.py
index 8bcced273..8e0dcf040 100644
--- a/src/bloqade/ir/control/waveform.py
+++ b/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
diff --git a/src/bloqade/ir/location/bravais.py b/src/bloqade/ir/location/bravais.py
index de37c7ed7..771f9d6ba 100644
--- a/src/bloqade/ir/location/bravais.py
+++ b/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
diff --git a/src/bloqade/ir/location/location.py b/src/bloqade/ir/location/location.py
index b903b22b0..12bffe6e7 100644
--- a/src/bloqade/ir/location/location.py
+++ b/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(
diff --git a/src/bloqade/ir/routine/base.py b/src/bloqade/ir/routine/base.py
index ead1fb8aa..f1c2f7e1b 100644
--- a/src/bloqade/ir/routine/base.py
+++ b/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:
diff --git a/src/bloqade/ir/routine/bloqade.py b/src/bloqade/ir/routine/bloqade.py
index 5c262eded..56d631701 100644
--- a/src/bloqade/ir/routine/bloqade.py
+++ b/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
diff --git a/src/bloqade/ir/routine/braket.py b/src/bloqade/ir/routine/braket.py
index adac18d15..577aedbad 100644
--- a/src/bloqade/ir/routine/braket.py
+++ b/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
diff --git a/src/bloqade/ir/routine/params.py b/src/bloqade/ir/routine/params.py
index 4f221915e..63e8227c3 100644
--- a/src/bloqade/ir/routine/params.py
+++ b/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
 
diff --git a/src/bloqade/ir/routine/quera.py b/src/bloqade/ir/routine/quera.py
index b1355df6b..c818caef5 100644
--- a/src/bloqade/ir/routine/quera.py
+++ b/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
diff --git a/src/bloqade/ir/scalar.py b/src/bloqade/ir/scalar.py
index ea876a1b9..be9b75226 100644
--- a/src/bloqade/ir/scalar.py
+++ b/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
diff --git a/src/bloqade/ir/tree_print.py b/src/bloqade/ir/tree_print.py
index 3d04337a4..b4868521d 100644
--- a/src/bloqade/ir/tree_print.py
+++ b/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
 
diff --git a/src/bloqade/submission/base.py b/src/bloqade/submission/base.py
index b8f575818..43384e632 100644
--- a/src/bloqade/submission/base.py
+++ b/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
 
diff --git a/src/bloqade/submission/braket.py b/src/bloqade/submission/braket.py
index 4298ca892..841c74976 100644
--- a/src/bloqade/submission/braket.py
+++ b/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
 
 
diff --git a/src/bloqade/submission/ir/braket.py b/src/bloqade/submission/ir/braket.py
index 20486209e..5a2077149 100644
--- a/src/bloqade/submission/ir/braket.py
+++ b/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
 
 
diff --git a/src/bloqade/submission/ir/capabilities.py b/src/bloqade/submission/ir/capabilities.py
index f95fc5f92..dbbdc566c 100644
--- a/src/bloqade/submission/ir/capabilities.py
+++ b/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
diff --git a/src/bloqade/submission/ir/parallel.py b/src/bloqade/submission/ir/parallel.py
index d334e8489..a511e0930 100644
--- a/src/bloqade/submission/ir/parallel.py
+++ b/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
diff --git a/src/bloqade/submission/ir/task_results.py b/src/bloqade/submission/ir/task_results.py
index 68ae250e7..7977beebc 100644
--- a/src/bloqade/submission/ir/task_results.py
+++ b/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
 
diff --git a/src/bloqade/submission/ir/task_specification.py b/src/bloqade/submission/ir/task_specification.py
index cbdc571be..f02524dc0 100644
--- a/src/bloqade/submission/ir/task_specification.py
+++ b/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
diff --git a/src/bloqade/submission/quera.py b/src/bloqade/submission/quera.py
index 6404016e7..b7802c511 100644
--- a/src/bloqade/submission/quera.py
+++ b/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,
diff --git a/src/bloqade/task/base.py b/src/bloqade/task/base.py
index 90dc77e17..275f74c01 100644
--- a/src/bloqade/task/base.py
+++ b/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
diff --git a/tests/test_builder.py b/tests/test_builder.py
index 95be58122..f9dd37713 100644
--- a/tests/test_builder.py
+++ b/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])
 
 
diff --git a/tests/test_scalar.py b/tests/test_scalar.py
index 4e9e22269..29a8515ad 100644
--- a/tests/test_scalar.py
+++ b/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