Adds attributes to print #15

UMD-ARLIS · Feb 26, 2022 · 111f813 · 111f813
1 parent 1c8a592
commit 111f813
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 3 deletions.
diff --git a/distill/segmentation/segment.py b/distill/segmentation/segment.py
@@ -58,8 +58,16 @@ def __init__(self, segment_name="", start_end_val=None, num_logs=0, uids=[]):
  def __str__(self):
  start = self.start_end_val[0]
  end = self.start_end_val[1]
- return "Segment: name=" + self.segment_name + ", num_logs=" + str(self.num_logs) + \
- ", start=" + str(start) + ", end=" + str(end) + ", type=" + str(self.segment_type)
+ variables = vars(self)
+ final_str = "Segment:"
+ for var in variables:
+ if var != "uids":
+ if var == "start_end_val":
+ final_str += " start" + "=" + str(start) + ","
+ final_str += " end" + "=" + str(end) + ","
+ else:
+ final_str += " " + str(var) + "=" + str(variables[var]) + ","
+ return final_str[:-1]
 
  def get_segment_name(self):
  """

diff --git a/tests/test_segment.py b/tests/test_segment.py
@@ -30,13 +30,28 @@
 ########################
 # SEGMENT OBJECT TESTS #
 ########################
+
 def test_segment_constructor():
  segment = distill.Segment()
  assert segment.get_segment_name() == ""
  assert segment.get_num_logs() == 0
  assert segment.get_start_end_val() is None
  assert segment.get_segment_uids() == []
 
+def test_segment_string():
+ segment = distill.Segment()
+ segment.segment_name = "segment_name"
+ segment.start_end_val = (1, 2)
+ segment.segment_type = distill.Segment_Type.CREATE
+
+ assert str(segment) == "Segment: segment_name=segment_name, start=1, end=2, num_logs=0, " \
+ "generate_field_name=None, generate_matched_values=None, segment_type=Segment_Type.CREATE"
+
+ segment.test = "test_attribute"
+ assert str(segment) == "Segment: segment_name=segment_name, start=1, end=2, num_logs=0, " \
+ "generate_field_name=None, generate_matched_values=None, segment_type=Segment_Type.CREATE," \
+ " test=test_attribute"
+
 def test_getters():
  data = testing_utils.setup("./data/sample_data.json", "integer")
  sorted_data = data[0]

diff --git a/tests/test_segments.py b/tests/test_segments.py
@@ -384,5 +384,6 @@ def test_str():
  segments = distill.Segments([segment])
 
  assert str(segments) == "Segments: [\n" \
- "Segment: name=segment_name, num_logs=5, start=1, end=2, type=None\n" \
+ "Segment: segment_name=segment_name, start=1, end=2, num_logs=5, " \
+ "generate_field_name=None, generate_matched_values=None, segment_type=None\n" \
  "]"