Skip to content

Commit

Permalink
Adding LocationEncoder to store Locations as JSON
Browse files Browse the repository at this point in the history
  • Loading branch information
gumyr committed Nov 26, 2023
1 parent fb6f474 commit 89fda66
Show file tree
Hide file tree
Showing 4 changed files with 94 additions and 5 deletions.
1 change: 1 addition & 0 deletions docs/direct_api_reference.rst
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ CAD objects described in the following section are frequently of these types.
:special-members: __copy__,__deepcopy__
.. autoclass:: Location
:special-members: __copy__,__deepcopy__, __mul__, __pow__, __eq__, __neg__
.. autoclass:: LocationEncoder
.. autoclass:: Pos
.. autoclass:: Rot
.. autoclass:: Matrix
Expand Down
1 change: 1 addition & 0 deletions src/build123d/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,7 @@
"Plane",
"Compound",
"Location",
"LocationEncoder",
"Joint",
"RigidJoint",
"RevoluteJoint",
Expand Down
49 changes: 46 additions & 3 deletions src/build123d/geometry.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
# too-many-arguments, too-many-locals, too-many-public-methods,
# too-many-statements, too-many-instance-attributes, too-many-branches
import copy
import json
import logging
from math import degrees, pi, radians
from typing import Any, Iterable, List, Optional, Sequence, Tuple, Union, overload
Expand Down Expand Up @@ -1092,7 +1093,7 @@ def __init__(self, *args):
elif len(args) == 1:
translation = args[0]

if isinstance(translation, (Vector, tuple)):
if isinstance(translation, (Vector, Iterable)):
transform.SetTranslationPart(Vector(translation).wrapped)
elif isinstance(translation, Plane):
coordinate_system = gp_Ax3(
Expand All @@ -1114,8 +1115,8 @@ def __init__(self, *args):
raise TypeError("Unexpected parameters")

elif len(args) == 2:
if isinstance(args[0], (Vector, tuple)):
if isinstance(args[1], (Vector, tuple)):
if isinstance(args[0], (Vector, Iterable)):
if isinstance(args[1], (Vector, Iterable)):
rotation = [radians(a) for a in args[1]]
quaternion = gp_Quaternion()
quaternion.SetEulerAngles(
Expand Down Expand Up @@ -1248,6 +1249,48 @@ def __str__(self):
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):
"""Subclass of Location used only for object rotation
Expand Down
48 changes: 46 additions & 2 deletions tests/test_direct_api.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
# system modules
import copy
import json
import math
import os
import platform
Expand Down Expand Up @@ -53,6 +54,7 @@
BoundBox,
Color,
Location,
LocationEncoder,
Matrix,
Pos,
Rot,
Expand Down Expand Up @@ -1393,6 +1395,12 @@ def test_location(self):
T = loc0.wrapped.Transformation().TranslationPart()
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
loc1 = Location(Vector(0, 0, 1))

Expand Down Expand Up @@ -1451,8 +1459,6 @@ def test_location(self):
self.assertAlmostEqual(30, angle7)

# Test error handling on creation
with self.assertRaises(TypeError):
Location([0, 0, 1])
with self.assertRaises(TypeError):
Location("xy_plane")

Expand Down Expand Up @@ -1557,6 +1563,44 @@ def test_neg(self):
self.assertVectorAlmostEquals(n_loc.position, (1, 2, 3), 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):
def test_matrix_creation_and_access(self):
Expand Down

0 comments on commit 89fda66

Please sign in to comment.