refactor: use match statements where idiomatic

pyproject.toml

@@ -14,7 +14,7 @@ requires = ['poetry-core']
 build-backend = 'poetry.core.masonry.api'
 
 [tool.poetry]
-version = '0.23.5rc1'
+version = '0.23.5rc2'
 packages = [{include = 'opvious', from = 'src'}]
 
 [tool.poetry.dependencies]
@@ -118,11 +118,7 @@ select = [
 ignore = [
   "ANN003", "ANN401",
   "D100", "D102", "D103", "D104", "D105", "D107", "D415",
-  "PLC0105",
-  "PLR2004",
-  "PLR0912",
-  "PLR0913",
-  "PLW2901",
+  "PLC0105", "PLR2004", "PLR0912", "PLR0913", "PLW2901",
   "RUF022",
   "TD002", "TD003",
 ]

src/opvious/client/handlers.py

@@ -4,7 +4,7 @@
 import dataclasses
 import json
 import logging
-from typing import Any, BinaryIO
+from typing import Any, BinaryIO, Self
 
 import backoff
 
@@ -136,17 +136,17 @@ def from_environment(
 
         token = ClientSetting.TOKEN.read(env).strip()
         authorization = authorization_header(token) if token else None
-        executor_key = ClientSetting.EXECUTOR.read(env)
-        if not executor_key:
-            executor = default_executor(endpoint, authorization)
-        elif executor_key == "aiohttp":
-            executor = aiohttp_executor(endpoint, authorization)
-        elif executor_key == "pyodide":
-            executor = pyodide_executor(endpoint, authorization)
-        elif executor_key == "urllib":
-            executor = urllib_executor(endpoint, authorization)
-        else:
-            raise ValueError(f"Unknown executor: {executor_key}")
+        match ClientSetting.EXECUTOR.read(env):
+            case "":
+                executor = default_executor(endpoint, authorization)
+            case "aiohttp":
+                executor = aiohttp_executor(endpoint, authorization)
+            case "pyodide":
+                executor = pyodide_executor(endpoint, authorization)
+            case "urllib":
+                executor = urllib_executor(endpoint, authorization)
+            case key:
+                raise ValueError(f"Unknown executor: {key}")
         return Client(executor)
 
     @property
@@ -344,7 +344,7 @@ async def summarize_problem(self, problem: Problem) -> ProblemSummary:
             return problem_summary_from_json(res.json_data())
 
     async def format_problem(
-        self, problem: Problem, include_line_comments: bool=False
+        self, problem: Problem, include_line_comments: bool = False
     ) -> str:
         r"""Returns the problem's annotated representation in `LP format`_
 
@@ -397,8 +397,8 @@ async def format_problem(
     async def solve(
         self,
         problem: Problem,
-        assert_feasible: bool=False,
-        prefer_streaming: bool=True,
+        assert_feasible: bool = False,
+        prefer_streaming: bool = True,
     ) -> Solution:
         """Solves an optimization problem remotely
 
@@ -448,59 +448,20 @@ async def solve(
         """
         prepared = await self._prepare_problem(problem)
         if prefer_streaming and self._executor.supports_streaming:
-            problem_summary = None
-            response_json = None
             async with self._executor.execute(
                 result_type=JsonSeqExecutorResult,
                 url="/solve",
                 method="POST",
                 json_data=json_dict(problem=prepared.data),
             ) as res:
-                async for data in res.json_seq_data():
-                    kind = data["kind"]
-                    if kind == "reifying":
-                        progress = data["progress"]
-                        if progress["kind"] == "constraint":
-                            summary = progress["summary"]
-                            _logger.debug(
-                                "Reified constraint %r. [columns=%s, rows=%s]",
-                                summary["label"],
-                                summary["columnCount"],
-                                summary["rowCount"],
-                            )
-                    elif kind == "reified":
-                        problem_summary = problem_summary_from_json(
-                            data["summary"]
-                        )
-                        _logger.info(
-                            "Solving problem... [columns=%s, rows=%s]",
-                            problem_summary.column_count,
-                            problem_summary.row_count,
-                        )
-                    elif kind == "solving":
-                        log_progress(_logger, data["progress"])
-                    elif kind == "denormalized":
-                        pass  # TODO: Output solution summary
-                    elif kind == "solved":
-                        _logger.debug("Downloaded outputs.")
-                        response_json = data
-                    elif kind == "error":
-                        message = "Solve failed"
-                        if res.trace:
-                            message += f" ({res.trace})"
-                        message += f": {data['error']['message']}"
-                        raise Exception(message)
-                    else:
-                        raise Exception(
-                            f"Unexpected response: {json.dumps(data)}"
-                        )
-            if not problem_summary or not response_json:
+                sink = await _SolveStreamSink.consume(res)
+            if not sink.problem_summary or not sink.response_json:
                 raise Exception("Streaming solve terminated early")
             solution = solution_from_json(
                 outline=prepared.outline,
                 inputs=prepared.inputs,
-                response_json=response_json,
-                problem_summary=problem_summary,
+                response_json=sink.response_json,
+                problem_summary=sink.problem_summary,
             )
         else:
             async with self._executor.execute(
@@ -671,7 +632,7 @@ async def _track_solve(self, uuid: Uuid) -> SolveOutcome | None:
     async def wait_for_solve_outcome(
         self,
         uuid: Uuid,
-        assert_feasible: bool=False,
+        assert_feasible: bool = False,
     ) -> SolveOutcome:
         """Waits for the solve to complete and returns its outcome
 
@@ -852,6 +813,56 @@ async def _next_page() -> list[QueuedSolve]:
             limit -= len(solves)
 
 
+class _SolveStreamSink:
+    def __init__(self, result: JsonSeqExecutorResult) -> None:
+        self._result = result
+        self.response_json: Json = None
+        self.problem_summary: ProblemSummary | None = None
+
+    @classmethod
+    async def consume(cls, result: JsonSeqExecutorResult) -> Self:
+        sink = cls(result)
+        async for data in result.json_seq_data():
+            sink._on_data(data)
+        return sink
+
+    def _on_data(self, data: Json) -> None:
+        match data["kind"]:
+            case "reifying":
+                progress = data["progress"]
+                if progress["kind"] == "constraint":
+                    summary = progress["summary"]
+                    _logger.debug(
+                        "Reified constraint %r. [columns=%s, rows=%s]",
+                        summary["label"],
+                        summary["columnCount"],
+                        summary["rowCount"],
+                    )
+            case "reified":
+                summary = problem_summary_from_json(data["summary"])
+                _logger.info(
+                    "Solving problem... [columns=%s, rows=%s]",
+                    summary.column_count,
+                    summary.row_count,
+                )
+                self.problem_summary = summary
+            case "solving":
+                log_progress(_logger, data["progress"])
+            case "denormalized":
+                pass  # TODO: Output solution summary
+            case "solved":
+                _logger.debug("Downloaded outputs.")
+                self.response_json = data
+            case "error":
+                message = "Solve failed"
+                if trace := self._result.trace:
+                    message += f" ({trace})"
+                message += f": {data['error']['message']}"
+                raise Exception(message)
+            case _:
+                raise Exception(f"Unexpected response: {json.dumps(data)}")
+
+
 @dataclasses.dataclass(frozen=True)
 class _PreparedProblem:
     data: Json

src/opvious/common.py

@@ -77,11 +77,13 @@ def encode_extended_float(val: ExtendedFloat) -> Json:
 
 
 def decode_extended_float(val: ExtendedFloat) -> Json:
-    if val == "Infinity":
-        return math.inf
-    elif val == "-Infinity":
-        return -math.inf
-    return val
+    match val:
+        case "Infinity":
+            return math.inf
+        case "-Infinity":
+            return -math.inf
+        case _:
+            return val
 
 
 def json_dict(**kwargs) -> Json:

src/opvious/data/outcomes.py

@@ -52,7 +52,7 @@ def failed_outcome_from_graphql(data: Any) -> FailedOutcome:
 class FeasibleOutcome:
     """A solution was found"""
 
-    optimal: bool
+    is_optimal: bool
     """Whether this solution was optimal (within gap thresholds)"""
 
     objective_value: Value | None
@@ -91,7 +91,7 @@ def solve_outcome_from_graphql(data: Any) -> SolveOutcome:
             return UnboundedOutcome()
         case "FEASIBLE" | "OPTIMAL":
             return FeasibleOutcome(
-                optimal=data["status"] == "OPTIMAL",
+                is_optimal=status == "OPTIMAL",
                 objective_value=data.get("objectiveValue"),
                 relative_gap=data.get("relativeGap"),
             )
@@ -104,7 +104,7 @@ def solve_outcome_status(outcome: SolveOutcome) -> SolveStatus:
     if isinstance(outcome, AbortedOutcome):
         return "ABORTED"
     if isinstance(outcome, FeasibleOutcome):
-        return "OPTIMAL" if outcome.optimal else "FEASIBLE"
+        return "OPTIMAL" if outcome.is_optimal else "FEASIBLE"
     if isinstance(outcome, InfeasibleOutcome):
         return "INFEASIBLE"
     if isinstance(outcome, UnboundedOutcome):

src/opvious/data/solves.py

@@ -335,19 +335,19 @@ def solution_from_json(
     problem_summary: ProblemSummary | None = None,
 ) -> Solution:
     outcome_json = response_json["outcome"]
-    status = outcome_json["status"]
-    if status == "INFEASIBLE":
-        outcome = cast(SolveOutcome, InfeasibleOutcome())
-    elif status == "UNBOUNDED":
-        outcome = UnboundedOutcome()
-    elif status == "ABORTED":
-        outcome = AbortedOutcome()
-    else:
-        outcome = FeasibleOutcome(
-            optimal=status == "OPTIMAL",
-            objective_value=outcome_json.get("objectiveValue"),
-            relative_gap=outcome_json.get("relativeGap"),
-        )
+    match status := outcome_json["status"]:
+        case "INFEASIBLE":
+            outcome = cast(SolveOutcome, InfeasibleOutcome())
+        case "UNBOUNDED":
+            outcome = UnboundedOutcome()
+        case "ABORTED":
+            outcome = AbortedOutcome()
+        case _:
+            outcome = FeasibleOutcome(
+                is_optimal=status == "OPTIMAL",
+                objective_value=outcome_json.get("objectiveValue"),
+                relative_gap=outcome_json.get("relativeGap"),
+            )
     outputs = None
     if isinstance(outcome, FeasibleOutcome):
         outputs = _outputs_from_json(

src/opvious/modeling/ast.py

@@ -236,25 +236,26 @@ def render(self, precedence: int=0) -> str:
         left_inner, right_inner, outer = _binary_operator_precedences[op]
         left = self.left_expression.render(left_inner)
         right = self.right_expression.render(right_inner)
-        if op == "mul":
-            if is_literal(self.left_expression, -1):
-                rendered = f"{{-{right}}}"
-            elif is_literal(self.right_expression, -1):
-                rendered = f"{{-{left}}}"
-            else:
-                rendered = f"{left} {right}"
-        elif op == "add":
-            rendered = f"{left} + {right}"
-        elif op == "mod":
-            rendered = f"{left} \\bmod {right}"
-        elif op == "sub":
-            rendered = f"{left} - {right}"
-        elif op == "div":
-            rendered = f"\\frac{{{left}}}{{{right}}}"
-        elif op == "pow":
-            rendered = f"\\left({left}\\right)^{{{right}}}"
-        else:
-            raise Exception(f"Unexpected operator: {op}")
+        match op:
+            case "mul":
+                if is_literal(self.left_expression, -1):
+                    rendered = f"{{-{right}}}"
+                elif is_literal(self.right_expression, -1):
+                    rendered = f"{{-{left}}}"
+                else:
+                    rendered = f"{left} {right}"
+            case "add":
+                rendered = f"{left} + {right}"
+            case "mod":
+                rendered = f"{left} \\bmod {right}"
+            case "sub":
+                rendered = f"{left} - {right}"
+            case "div":
+                rendered = f"\\frac{{{left}}}{{{right}}}"
+            case "pow":
+                rendered = f"\\left({left}\\right)^{{{right}}}"
+            case _:
+                raise Exception(f"Unexpected operator: {op}")
         if outer < precedence:
             rendered = f"\\left({rendered}\\right)"
         return rendered

src/opvious/transformations.py

@@ -6,6 +6,7 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 import dataclasses
 import math
 from typing import Literal
@@ -74,7 +75,7 @@ class PinVariables(ProblemTransformation):
     :class:`.RelaxConstraints`.
     """
 
-    labels: list[Label] = dataclasses.field(default_factory=lambda: [])
+    labels: Sequence[Label] = dataclasses.field(default_factory=lambda: [])
     """The labels of the variables to pin
 
     If empty, all variables will be pinned.
@@ -125,7 +126,7 @@ class RelaxConstraints(ProblemTransformation):
     aggregate slack violation (see :ref:`Detecting infeasibilities`).
     """
 
-    labels: list[Label] = dataclasses.field(default_factory=lambda: [])
+    labels: Sequence[Label] = dataclasses.field(default_factory=lambda: [])
     """The labels of the constraints to relax
 
     If empty, all constraints will be relaxed.
@@ -173,7 +174,7 @@ async def register(self, context: ProblemTransformationContext) -> None:
 class DensifyVariables(ProblemTransformation):
     """A transformation which updates one or more variables to be continuous"""
 
-    labels: list[Label] = dataclasses.field(default_factory=lambda: [])
+    labels: Sequence[Label] = dataclasses.field(default_factory=lambda: [])
     """The labels of the variables to densify
 
     If empty, all integral variables will be densified.
@@ -196,7 +197,7 @@ class OmitConstraints(ProblemTransformation):
     constraint or objective will automatically be dropped.
     """
 
-    labels: list[Label] = dataclasses.field(default_factory=lambda: [])
+    labels: Sequence[Label] = dataclasses.field(default_factory=lambda: [])
     """The labels of the constraints to drop
 
     If empty, all constraints will be dropped.
@@ -220,7 +221,7 @@ class OmitObjectives(ProblemTransformation):
     be dropped.
     """
 
-    labels: list[Label] = dataclasses.field(default_factory=lambda: [])
+    labels: Sequence[Label] = dataclasses.field(default_factory=lambda: [])
     """The labels of the objectives to drop
 
     If empty, all objectives will be dropped.