From 08474a1e3147400858cd15c39fd4411160a618ba Mon Sep 17 00:00:00 2001
From: David Vo <auscompgeek@users.noreply.github.com>
Date: Thu, 26 Jan 2023 23:17:46 +1100
Subject: [PATCH] PhotonPoseEstimator: Move an lvalue reference

---
 gen/PhotonPoseEstimator.yml  |  7 ++++++
 gen/RobotPoseEstimator.yml   |  1 +
 photonvision/__init__.py     |  2 ++
 pyproject.toml               |  2 +-
 tests/test_pose_estimator.py | 46 ++++++++++++++++++++++++++++++++++++
 5 files changed, 57 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_pose_estimator.py

diff --git a/gen/PhotonPoseEstimator.yml b/gen/PhotonPoseEstimator.yml
index 9e75061..06df092 100644
--- a/gen/PhotonPoseEstimator.yml
+++ b/gen/PhotonPoseEstimator.yml
@@ -17,6 +17,12 @@ classes:
   PhotonPoseEstimator:
     methods:
       PhotonPoseEstimator:
+        cpp_code: |
+          [](frc::AprilTagFieldLayout aprilTags, photonlib::PoseStrategy strategy,
+             PhotonCamera& camera, frc::Transform3d robotToCamera) {
+            return std::make_unique<PhotonPoseEstimator>(
+              aprilTags, strategy, std::move(camera), robotToCamera);
+          }
       GetFieldLayout:
       GetPoseStrategy:
       SetPoseStrategy:
@@ -25,3 +31,4 @@ classes:
       SetLastPose:
       Update:
       GetCamera:
+        return_value_policy: reference_internal
diff --git a/gen/RobotPoseEstimator.yml b/gen/RobotPoseEstimator.yml
index f844f20..e325428 100644
--- a/gen/RobotPoseEstimator.yml
+++ b/gen/RobotPoseEstimator.yml
@@ -4,6 +4,7 @@ extra_includes:
 
 enums:
   PoseStrategy:
+    ignore: true  # duplicated in PhotonPoseEstimator.h
 classes:
   RobotPoseEstimator:
     force_type_casters:
diff --git a/photonvision/__init__.py b/photonvision/__init__.py
index 886ef3e..2486f5e 100644
--- a/photonvision/__init__.py
+++ b/photonvision/__init__.py
@@ -6,6 +6,7 @@
     Packet,
     PhotonCamera,
     PhotonPipelineResult,
+    PhotonPoseEstimator,
     PhotonTrackedTarget,
     PhotonUtils,
     PoseStrategy,
@@ -21,6 +22,7 @@
     "Packet",
     "PhotonCamera",
     "PhotonPipelineResult",
+    "PhotonPoseEstimator",
     "PhotonTrackedTarget",
     "PhotonUtils",
     "PoseStrategy",
diff --git a/pyproject.toml b/pyproject.toml
index f8ed508..b9b9e9e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -57,7 +57,7 @@ generation_data = "gen"
 Packet = "photonlib/Packet.h"
 PhotonCamera = "photonlib/PhotonCamera.h"
 PhotonPipelineResult = "photonlib/PhotonPipelineResult.h"
-#PhotonPoseEstimator = "photonlib/PhotonPoseEstimator.h"
+PhotonPoseEstimator = "photonlib/PhotonPoseEstimator.h"
 PhotonTargetSortMode = "photonlib/PhotonTargetSortMode.h"
 PhotonTrackedTarget = "photonlib/PhotonTrackedTarget.h"
 PhotonUtils = "photonlib/PhotonUtils.h"
diff --git a/tests/test_pose_estimator.py b/tests/test_pose_estimator.py
new file mode 100644
index 0000000..b8f84eb
--- /dev/null
+++ b/tests/test_pose_estimator.py
@@ -0,0 +1,46 @@
+import robotpy_apriltag
+from wpimath.geometry import Transform3d
+
+import photonvision
+
+
+def test_photon_pose_estimator_ctor():
+    camera = photonvision.PhotonCamera("test_ctor")
+    assert camera._path
+    estimator = photonvision.PhotonPoseEstimator(
+        robotpy_apriltag.loadAprilTagLayoutField(
+            robotpy_apriltag.AprilTagField.k2023ChargedUp
+        ),
+        photonvision.PoseStrategy.AVERAGE_BEST_TARGETS,
+        camera,
+        Transform3d(),
+    )
+    assert estimator
+    assert not camera._path
+
+
+def test_photon_pose_estimator_update():
+    estimator = photonvision.PhotonPoseEstimator(
+        robotpy_apriltag.loadAprilTagLayoutField(
+            robotpy_apriltag.AprilTagField.k2023ChargedUp
+        ),
+        photonvision.PoseStrategy.AVERAGE_BEST_TARGETS,
+        photonvision.PhotonCamera("test_update"),
+        Transform3d(),
+    )
+    assert estimator.update() is None
+
+
+def test_photon_pose_estimator_get_cam():
+    camera = photonvision.PhotonCamera("test_get_cam")
+    estimator = photonvision.PhotonPoseEstimator(
+        robotpy_apriltag.loadAprilTagLayoutField(
+            robotpy_apriltag.AprilTagField.k2023ChargedUp
+        ),
+        photonvision.PoseStrategy.AVERAGE_BEST_TARGETS,
+        camera,
+        Transform3d(),
+    )
+    inner_cam = estimator.getCamera()
+    assert inner_cam is not None
+    assert inner_cam is not camera