From 77965db122046b205b04f3f24026ba8281c877aa Mon Sep 17 00:00:00 2001
From: Yi-Ting Lee <yitingle@amazon.com>
Date: Fri, 1 Jul 2022 09:07:32 -0700
Subject: [PATCH] feature: add __str__ methods to queryLineage output classes

---
 src/sagemaker/lineage/query.py | 51 ++++++++++++++++++++++++++++++++++
 1 file changed, 51 insertions(+)

diff --git a/src/sagemaker/lineage/query.py b/src/sagemaker/lineage/query.py
index a54331c39a..72bde00a1a 100644
--- a/src/sagemaker/lineage/query.py
+++ b/src/sagemaker/lineage/query.py
@@ -92,6 +92,18 @@ def __eq__(self, other):
             and self.destination_arn == other.destination_arn
         )
 
+    def __str__(self):
+        """Define string representation of ``Edge``.
+
+        Format:
+            {
+                'source_arn': 'string', 'destination_arn': 'string', 
+                'association_type': 'string'
+            }
+        
+        """
+        return (str(self.__dict__))
+
 
 class Vertex:
     """A vertex for a lineage graph."""
@@ -130,6 +142,19 @@ def __eq__(self, other):
             and self.lineage_source == other.lineage_source
         )
 
+    def __str__(self):
+        """Define string representation of ``Vertex``.
+
+        Format:
+            {
+                'arn': 'string', 'lineage_entity': 'string', 
+                'lineage_source': 'string', 
+                '_session': <sagemaker.session.Session object>
+            }
+        
+        """
+        return (str(self.__dict__))
+
     def to_lineage_object(self):
         """Convert the ``Vertex`` object to its corresponding lineage object.
 
@@ -199,6 +224,32 @@ def __init__(
         if vertices is not None:
             self.vertices = vertices
 
+    def __str__(self):
+        """Define string representation of ``LineageQueryResult``.
+        
+        Format:
+        {
+            'edges':[
+                {
+                    'source_arn': 'string', 'destination_arn': 'string', 
+                    'association_type': 'string'
+                },
+                ...
+            ]
+            'vertices':[
+                {
+                    'arn': 'string', 'lineage_entity': 'string', 
+                    'lineage_source': 'string', 
+                    '_session': <sagemaker.session.Session object>
+                },
+                ...
+            ]
+        }
+        
+        """
+        result_dict = vars(self)
+        return (str({k: [vars(val) for val in v] for k, v in result_dict.items()}))
+
 
 class LineageFilter(object):
     """A filter used in a lineage query."""