Adding LocationEncoder to store Locations as JSON

This commit is contained in:
gumyr 2023-11-25 19:13:32 -05:00
parent fb6f47471f
commit 89fda66873
4 changed files with 94 additions and 5 deletions

View file

@ -30,6 +30,7 @@ CAD objects described in the following section are frequently of these types.
:special-members: __copy__,__deepcopy__ :special-members: __copy__,__deepcopy__
.. autoclass:: Location .. autoclass:: Location
:special-members: __copy__,__deepcopy__, __mul__, __pow__, __eq__, __neg__ :special-members: __copy__,__deepcopy__, __mul__, __pow__, __eq__, __neg__
.. autoclass:: LocationEncoder
.. autoclass:: Pos .. autoclass:: Pos
.. autoclass:: Rot .. autoclass:: Rot
.. autoclass:: Matrix .. autoclass:: Matrix

View file

@ -138,6 +138,7 @@ __all__ = [
"Plane", "Plane",
"Compound", "Compound",
"Location", "Location",
"LocationEncoder",
"Joint", "Joint",
"RigidJoint", "RigidJoint",
"RevoluteJoint", "RevoluteJoint",

View file

@ -34,6 +34,7 @@ from __future__ import annotations
# too-many-arguments, too-many-locals, too-many-public-methods, # too-many-arguments, too-many-locals, too-many-public-methods,
# too-many-statements, too-many-instance-attributes, too-many-branches # too-many-statements, too-many-instance-attributes, too-many-branches
import copy import copy
import json
import logging import logging
from math import degrees, pi, radians from math import degrees, pi, radians
from typing import Any, Iterable, List, Optional, Sequence, Tuple, Union, overload from typing import Any, Iterable, List, Optional, Sequence, Tuple, Union, overload
@ -1092,7 +1093,7 @@ class Location:
elif len(args) == 1: elif len(args) == 1:
translation = args[0] translation = args[0]
if isinstance(translation, (Vector, tuple)): if isinstance(translation, (Vector, Iterable)):
transform.SetTranslationPart(Vector(translation).wrapped) transform.SetTranslationPart(Vector(translation).wrapped)
elif isinstance(translation, Plane): elif isinstance(translation, Plane):
coordinate_system = gp_Ax3( coordinate_system = gp_Ax3(
@ -1114,8 +1115,8 @@ class Location:
raise TypeError("Unexpected parameters") raise TypeError("Unexpected parameters")
elif len(args) == 2: elif len(args) == 2:
if isinstance(args[0], (Vector, tuple)): if isinstance(args[0], (Vector, Iterable)):
if isinstance(args[1], (Vector, tuple)): if isinstance(args[1], (Vector, Iterable)):
rotation = [radians(a) for a in args[1]] rotation = [radians(a) for a in args[1]]
quaternion = gp_Quaternion() quaternion = gp_Quaternion()
quaternion.SetEulerAngles( quaternion.SetEulerAngles(
@ -1248,6 +1249,48 @@ class Location:
return f"Location: (position=({position_str}), orientation=({orientation_str}))" return f"Location: (position=({position_str}), orientation=({orientation_str}))"
class LocationEncoder(json.JSONEncoder):
"""Custom JSON Encoder for Location values
Example:
.. code::
data_dict = {
"part1": {
"joint_one": Location((1, 2, 3), (4, 5, 6)),
"joint_two": Location((7, 8, 9), (10, 11, 12)),
},
"part2": {
"joint_one": Location((13, 14, 15), (16, 17, 18)),
"joint_two": Location((19, 20, 21), (22, 23, 24)),
},
}
json_object = json.dumps(data_dict, indent=4, cls=LocationEncoder)
with open("sample.json", "w") as outfile:
outfile.write(json_object)
with open("sample.json", "r") as infile:
copy_data_dict = json.load(infile, object_hook=LocationEncoder.location_hook)
"""
def default(self, loc: Location) -> dict:
"""Return a serializable object"""
if not isinstance(loc, Location):
raise TypeError("Only applies to Location objects")
return {"Location": loc.to_tuple()}
def location_hook(obj) -> dict:
"""Convert Locations loaded from json to Location objects
Example:
read_json = json.load(infile, object_hook=LocationEncoder.location_hook)
"""
if "Location" in obj:
obj = Location(*[[float(f) for f in v] for v in obj["Location"]])
return obj
class Rotation(Location): class Rotation(Location):
"""Subclass of Location used only for object rotation """Subclass of Location used only for object rotation

View file

@ -1,5 +1,6 @@
# system modules # system modules
import copy import copy
import json
import math import math
import os import os
import platform import platform
@ -53,6 +54,7 @@ from build123d.geometry import (
BoundBox, BoundBox,
Color, Color,
Location, Location,
LocationEncoder,
Matrix, Matrix,
Pos, Pos,
Rot, Rot,
@ -1393,6 +1395,12 @@ class TestLocation(DirectApiTestCase):
T = loc0.wrapped.Transformation().TranslationPart() T = loc0.wrapped.Transformation().TranslationPart()
self.assertTupleAlmostEquals((T.X(), T.Y(), T.Z()), (0, 0, 1), 6) self.assertTupleAlmostEquals((T.X(), T.Y(), T.Z()), (0, 0, 1), 6)
# List
loc0 = Location([0, 0, 1])
T = loc0.wrapped.Transformation().TranslationPart()
self.assertTupleAlmostEquals((T.X(), T.Y(), T.Z()), (0, 0, 1), 6)
# Vector # Vector
loc1 = Location(Vector(0, 0, 1)) loc1 = Location(Vector(0, 0, 1))
@ -1451,8 +1459,6 @@ class TestLocation(DirectApiTestCase):
self.assertAlmostEqual(30, angle7) self.assertAlmostEqual(30, angle7)
# Test error handling on creation # Test error handling on creation
with self.assertRaises(TypeError):
Location([0, 0, 1])
with self.assertRaises(TypeError): with self.assertRaises(TypeError):
Location("xy_plane") Location("xy_plane")
@ -1557,6 +1563,44 @@ class TestLocation(DirectApiTestCase):
self.assertVectorAlmostEquals(n_loc.position, (1, 2, 3), 5) self.assertVectorAlmostEquals(n_loc.position, (1, 2, 3), 5)
self.assertVectorAlmostEquals(n_loc.orientation, (180, -35, -127), 5) self.assertVectorAlmostEquals(n_loc.orientation, (180, -35, -127), 5)
def test_as_json(self):
data_dict = {
"part1": {
"joint_one": Location((1, 2, 3), (4, 5, 6)),
"joint_two": Location((7, 8, 9), (10, 11, 12)),
},
"part2": {
"joint_one": Location((13, 14, 15), (16, 17, 18)),
"joint_two": Location((19, 20, 21), (22, 23, 24)),
},
}
# Serializing json with custom Location encoder
json_object = json.dumps(data_dict, indent=4, cls=LocationEncoder)
# Writing to sample.json
with open("sample.json", "w") as outfile:
outfile.write(json_object)
# Reading from sample.json
with open("sample.json", "r") as infile:
read_json = json.load(infile, object_hook=LocationEncoder.location_hook)
# Validate locations
for key, value in read_json.items():
for k, v in value.items():
if key == "part1" and k == "joint_one":
self.assertVectorAlmostEquals(v.position, (1, 2, 3), 5)
elif key == "part1" and k == "joint_two":
self.assertVectorAlmostEquals(v.position, (7, 8, 9), 5)
elif key == "part2" and k == "joint_one":
self.assertVectorAlmostEquals(v.position, (13, 14, 15), 5)
elif key == "part2" and k == "joint_two":
self.assertVectorAlmostEquals(v.position, (19, 20, 21), 5)
else:
self.assertTrue(False)
os.remove("sample.json")
class TestMatrix(DirectApiTestCase): class TestMatrix(DirectApiTestCase):
def test_matrix_creation_and_access(self): def test_matrix_creation_and_access(self):