Skip to content

Commit

Permalink
Support pydantic v2 and get rid of plum warnings (#922)
Browse files Browse the repository at this point in the history
* Support pydantic v2 the easiest way possible and get rid of plum warnings

* whoops forgot lint

* bump version
  • Loading branch information
johnzl-777 authored Apr 2, 2024
1 parent e2561c6 commit 05a9f94
Show file tree
Hide file tree
Showing 33 changed files with 60 additions and 47 deletions.
2 changes: 1 addition & 1 deletion pdm.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 2 additions & 2 deletions pyproject.toml
Original file line number Diff line number Diff line change
@@ -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"},
Expand All @@ -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",
Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/compiler/analysis/common/scan_variables.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/compiler/analysis/python/waveform.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/compiler/codegen/hardware/lattice.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/compiler/codegen/hardware/piecewise_linear.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/compiler/passes/hardware/components.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/ir/analog_circuit.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/ir/control/field.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/ir/control/pulse.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/ir/control/sequence.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/ir/control/waveform.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/ir/location/bravais.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down
32 changes: 23 additions & 9 deletions src/bloqade/ir/location/location.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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,
Expand All @@ -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(
Expand Down Expand Up @@ -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(
Expand Down
4 changes: 2 additions & 2 deletions src/bloqade/ir/routine/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/ir/routine/bloqade.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/ir/routine/braket.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/ir/routine/params.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/ir/routine/quera.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down
4 changes: 2 additions & 2 deletions src/bloqade/ir/scalar.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/ir/tree_print.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/submission/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/submission/braket.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/submission/ir/braket.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
LocalField,
)
from typing import Tuple, Union, List
from pydantic import BaseModel
from pydantic.v1 import BaseModel
from decimal import Decimal


Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/submission/ir/capabilities.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/submission/ir/parallel.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/submission/ir/task_results.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/submission/ir/task_specification.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/submission/quera.py
Original file line number Diff line number Diff line change
@@ -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,
Expand Down
2 changes: 1 addition & 1 deletion src/bloqade/task/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
7 changes: 3 additions & 4 deletions tests/test_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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])


Expand Down
2 changes: 1 addition & 1 deletion tests/test_scalar.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down

0 comments on commit 05a9f94

Please sign in to comment.