From 89fda668736d2d8ed0b9049b272bf839e373cd47 Mon Sep 17 00:00:00 2001 From: gumyr Date: Sat, 25 Nov 2023 19:13:32 -0500 Subject: [PATCH] Adding LocationEncoder to store Locations as JSON --- docs/direct_api_reference.rst | 1 + src/build123d/__init__.py | 1 + src/build123d/geometry.py | 49 ++++++++++++++++++++++++++++++++--- tests/test_direct_api.py | 48 ++++++++++++++++++++++++++++++++-- 4 files changed, 94 insertions(+), 5 deletions(-) diff --git a/docs/direct_api_reference.rst b/docs/direct_api_reference.rst index 94bcf72..7243789 100644 --- a/docs/direct_api_reference.rst +++ b/docs/direct_api_reference.rst @@ -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 diff --git a/src/build123d/__init__.py b/src/build123d/__init__.py index 7babbab..5349c45 100644 --- a/src/build123d/__init__.py +++ b/src/build123d/__init__.py @@ -138,6 +138,7 @@ __all__ = [ "Plane", "Compound", "Location", + "LocationEncoder", "Joint", "RigidJoint", "RevoluteJoint", diff --git a/src/build123d/geometry.py b/src/build123d/geometry.py index b99b815..d8c356c 100644 --- a/src/build123d/geometry.py +++ b/src/build123d/geometry.py @@ -34,6 +34,7 @@ from __future__ import annotations # 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 @@ -1092,7 +1093,7 @@ class Location: 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( @@ -1114,8 +1115,8 @@ class Location: 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( @@ -1248,6 +1249,48 @@ class Location: 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 diff --git a/tests/test_direct_api.py b/tests/test_direct_api.py index b54889c..fecaa50 100644 --- a/tests/test_direct_api.py +++ b/tests/test_direct_api.py @@ -1,5 +1,6 @@ # system modules import copy +import json import math import os import platform @@ -53,6 +54,7 @@ from build123d.geometry import ( BoundBox, Color, Location, + LocationEncoder, Matrix, Pos, Rot, @@ -1393,6 +1395,12 @@ class TestLocation(DirectApiTestCase): 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)) @@ -1451,8 +1459,6 @@ class TestLocation(DirectApiTestCase): self.assertAlmostEqual(30, angle7) # Test error handling on creation - with self.assertRaises(TypeError): - Location([0, 0, 1]) with self.assertRaises(TypeError): Location("xy_plane") @@ -1557,6 +1563,44 @@ class TestLocation(DirectApiTestCase): 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):