Skip to content

Commit

Permalink
Update pydantic schema
Browse files Browse the repository at this point in the history
  • Loading branch information
@aborgna-q
aborgna-q committed Apr 11, 2024
1 parent 76c8ea4 commit 9a87e0a
Show file tree
Hide file tree
Showing 2 changed files with 76 additions and 165 deletions.
68 changes: 20 additions & 48 deletions hugr-py/src/hugr/serialization/ops.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -300,13 +300,6 @@ class LoadConstant(DataflowOp):
datatype: Type


class LeafOpBase(DataflowOp):
"""Simple operation that has only value inputs+outputs and (potentially) StateOrder
edges."""

op: Literal["LeafOp"] = "LeafOp"


class DFG(DataflowOp):
"""A simply nested dataflow graph."""

Expand Down Expand Up @@ -398,11 +391,11 @@ def insert_port_types(self, inputs: TypeRow, outputs: TypeRow) -> None:
# -----------------------------------------


class CustomOp(LeafOpBase):
class CustomOp(DataflowOp):
"""A user-defined operation that can be downcasted by the extensions that define
it."""

lop: Literal["CustomOp"] = "CustomOp"
op: Literal["CustomOp"] = "CustomOp"
extension: ExtensionId
op_name: str
signature: tys.FunctionType = Field(default_factory=tys.FunctionType.empty)
Expand All @@ -425,10 +418,10 @@ class Config:
}


class Noop(LeafOpBase):
class Noop(DataflowOp):
"""A no-op operation."""

lop: Literal["Noop"] = "Noop"
op: Literal["Noop"] = "Noop"
ty: Type

def insert_port_types(self, in_types: TypeRow, out_types: TypeRow) -> None:
Expand All @@ -438,10 +431,10 @@ def insert_port_types(self, in_types: TypeRow, out_types: TypeRow) -> None:
self.ty = in_types[0]


class MakeTuple(LeafOpBase):
class MakeTuple(DataflowOp):
"""An operation that packs all its inputs into a tuple."""

lop: Literal["MakeTuple"] = "MakeTuple"
op: Literal["MakeTuple"] = "MakeTuple"
tys: TypeRow = Field(default_factory=list)

def insert_port_types(self, in_types: TypeRow, out_types: TypeRow) -> None:
Expand All @@ -451,56 +444,30 @@ def insert_port_types(self, in_types: TypeRow, out_types: TypeRow) -> None:
self.tys = list(in_types)


class UnpackTuple(LeafOpBase):
class UnpackTuple(DataflowOp):
"""An operation that packs all its inputs into a tuple."""

lop: Literal["UnpackTuple"] = "UnpackTuple"
op: Literal["UnpackTuple"] = "UnpackTuple"
tys: TypeRow = Field(default_factory=list)

def insert_port_types(self, in_types: TypeRow, out_types: TypeRow) -> None:
self.tys = list(out_types)


class Tag(LeafOpBase):
class Tag(DataflowOp):
"""An operation that creates a tagged sum value from one of its variants."""

lop: Literal["Tag"] = "Tag"
op: Literal["Tag"] = "Tag"
tag: int # The variant to create.
variants: TypeRow # The variants of the sum type.


class TypeApply(LeafOpBase):
class Lift(DataflowOp):
"""Fixes some TypeParams of a polymorphic type by providing TypeArgs."""

lop: Literal["TypeApply"] = "TypeApply"
ta: "TypeApplication"


class TypeApplication(BaseModel):
"""Records details of an application of a PolyFuncType to some TypeArgs and the
result (a less-, but still potentially-, polymorphic type).
"""

input: PolyFuncType
args: list[tys.TypeTypeArg]
output: PolyFuncType

class Config:
# Needed to avoid random '\n's in the pydantic description
json_schema_extra = {
"description": (
"Records details of an application of a PolyFuncType to some TypeArgs "
"and the result (a less-, but still potentially-, polymorphic type)."
)
}


class LeafOp(RootModel):
"""A constant operation."""

root: CustomOp | Noop | MakeTuple | UnpackTuple | Tag | TypeApply = Field(
discriminator="lop"
)
op: Literal["Lift"] = "Lift"
type_row: TypeRow
new_extension: ExtensionId


class OpType(RootModel):
Expand All @@ -522,7 +489,12 @@ class OpType(RootModel):
| Call
| CallIndirect
| LoadConstant
| LeafOp
| CustomOp
| Noop
| MakeTuple
| UnpackTuple
| Tag
| Lift
| DFG
) = Field(discriminator="op")

Expand Down
173 changes: 56 additions & 117 deletions specification/schema/hugr_schema_v1.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -255,14 +255,9 @@
"$ref": "#/$defs/ExtensionSet"
},
"op": {
"const": "LeafOp",
"default": "LeafOp",
"title": "Op"
},
"lop": {
"const": "CustomOp",
"default": "CustomOp",
"title": "Lop"
"title": "Op"
},
"extension": {
"title": "Extension",
Expand Down Expand Up @@ -674,40 +669,40 @@
"title": "Input",
"type": "object"
},
"LeafOp": {
"description": "A constant operation.",
"discriminator": {
"mapping": {
"CustomOp": "#/$defs/CustomOp",
"MakeTuple": "#/$defs/MakeTuple",
"Noop": "#/$defs/Noop",
"Tag": "#/$defs/Tag",
"TypeApply": "#/$defs/TypeApply",
"UnpackTuple": "#/$defs/UnpackTuple"
},
"propertyName": "lop"
},
"oneOf": [
{
"$ref": "#/$defs/CustomOp"
},
{
"$ref": "#/$defs/Noop"
"Lift": {
"description": "Fixes some TypeParams of a polymorphic type by providing TypeArgs.",
"properties": {
"parent": {
"title": "Parent",
"type": "integer"
},
{
"$ref": "#/$defs/MakeTuple"
"input_extensions": {
"$ref": "#/$defs/ExtensionSet"
},
{
"$ref": "#/$defs/UnpackTuple"
"op": {
"const": "Lift",
"default": "Lift",
"title": "Op"
},
{
"$ref": "#/$defs/Tag"
"type_row": {
"items": {
"$ref": "#/$defs/Type"
},
"title": "Type Row",
"type": "array"
},
{
"$ref": "#/$defs/TypeApply"
"new_extension": {
"title": "New Extension",
"type": "string"
}
},
"required": [
"parent",
"type_row",
"new_extension"
],
"title": "LeafOp"
"title": "Lift",
"type": "object"
},
"ListParam": {
"properties": {
Expand Down Expand Up @@ -763,14 +758,9 @@
"$ref": "#/$defs/ExtensionSet"
},
"op": {
"const": "LeafOp",
"default": "LeafOp",
"title": "Op"
},
"lop": {
"const": "MakeTuple",
"default": "MakeTuple",
"title": "Lop"
"title": "Op"
},
"tys": {
"items": {
Expand Down Expand Up @@ -819,14 +809,9 @@
"$ref": "#/$defs/ExtensionSet"
},
"op": {
"const": "LeafOp",
"default": "LeafOp",
"title": "Op"
},
"lop": {
"const": "Noop",
"default": "Noop",
"title": "Lop"
"title": "Op"
},
"ty": {
"$ref": "#/$defs/Type"
Expand All @@ -849,17 +834,22 @@
"Case": "#/$defs/Case",
"Conditional": "#/$defs/Conditional",
"Const": "#/$defs/Const",
"CustomOp": "#/$defs/CustomOp",
"DFG": "#/$defs/DFG",
"DataflowBlock": "#/$defs/DataflowBlock",
"ExitBlock": "#/$defs/ExitBlock",
"FuncDecl": "#/$defs/FuncDecl",
"FuncDefn": "#/$defs/FuncDefn",
"Input": "#/$defs/Input",
"LeafOp": "#/$defs/LeafOp",
"Lift": "#/$defs/Lift",
"LoadConstant": "#/$defs/LoadConstant",
"MakeTuple": "#/$defs/MakeTuple",
"Module": "#/$defs/Module",
"Noop": "#/$defs/Noop",
"Output": "#/$defs/Output",
"TailLoop": "#/$defs/TailLoop"
"Tag": "#/$defs/Tag",
"TailLoop": "#/$defs/TailLoop",
"UnpackTuple": "#/$defs/UnpackTuple"
},
"propertyName": "op"
},
Expand Down Expand Up @@ -910,7 +900,22 @@
"$ref": "#/$defs/LoadConstant"
},
{
"$ref": "#/$defs/LeafOp"
"$ref": "#/$defs/CustomOp"
},
{
"$ref": "#/$defs/Noop"
},
{
"$ref": "#/$defs/MakeTuple"
},
{
"$ref": "#/$defs/UnpackTuple"
},
{
"$ref": "#/$defs/Tag"
},
{
"$ref": "#/$defs/Lift"
},
{
"$ref": "#/$defs/DFG"
Expand Down Expand Up @@ -1146,14 +1151,9 @@
"$ref": "#/$defs/ExtensionSet"
},
"op": {
"const": "LeafOp",
"default": "LeafOp",
"title": "Op"
},
"lop": {
"const": "Tag",
"default": "Tag",
"title": "Lop"
"title": "Op"
},
"tag": {
"title": "Tag",
Expand Down Expand Up @@ -1315,62 +1315,6 @@
],
"title": "Type"
},
"TypeApplication": {
"description": "Records details of an application of a PolyFuncType to some TypeArgs and the result (a less-, but still potentially-, polymorphic type).",
"properties": {
"input": {
"$ref": "#/$defs/PolyFuncType"
},
"args": {
"items": {
"$ref": "#/$defs/TypeTypeArg"
},
"title": "Args",
"type": "array"
},
"output": {
"$ref": "#/$defs/PolyFuncType"
}
},
"required": [
"input",
"args",
"output"
],
"title": "TypeApplication",
"type": "object"
},
"TypeApply": {
"description": "Fixes some TypeParams of a polymorphic type by providing TypeArgs.",
"properties": {
"parent": {
"title": "Parent",
"type": "integer"
},
"input_extensions": {
"$ref": "#/$defs/ExtensionSet"
},
"op": {
"const": "LeafOp",
"default": "LeafOp",
"title": "Op"
},
"lop": {
"const": "TypeApply",
"default": "TypeApply",
"title": "Lop"
},
"ta": {
"$ref": "#/$defs/TypeApplication"
}
},
"required": [
"parent",
"ta"
],
"title": "TypeApply",
"type": "object"
},
"TypeArg": {
"description": "A type argument.",
"discriminator": {
Expand Down Expand Up @@ -1523,14 +1467,9 @@
"$ref": "#/$defs/ExtensionSet"
},
"op": {
"const": "LeafOp",
"default": "LeafOp",
"title": "Op"
},
"lop": {
"const": "UnpackTuple",
"default": "UnpackTuple",
"title": "Lop"
"title": "Op"
},
"tys": {
"items": {
Expand Down

0 comments on commit 9a87e0a

Please sign in to comment.