From 52962402fe0353d63652db31f62277203c0016cd Mon Sep 17 00:00:00 2001
From: Jan Musial <jan.musial@intel.com>
Date: Tue, 7 Apr 2020 16:09:13 +0200
Subject: [PATCH] Extend Size class to enable arithmetics

Signed-off-by: Jan Musial <jan.musial@intel.com>
---
 tests/functional/pyocf/utils.py | 55 +++++++++++++++++++++++++++++++++
 1 file changed, 55 insertions(+)

diff --git a/tests/functional/pyocf/utils.py b/tests/functional/pyocf/utils.py
index d4ef423..2eb5fa3 100644
--- a/tests/functional/pyocf/utils.py
+++ b/tests/functional/pyocf/utils.py
@@ -148,6 +148,61 @@ class Size:
         else:
             return "{} TiB".format(self.TiB)
 
+    def __repr__(self):
+        return f"Size({self.bytes})"
+
+    def __eq__(self, other):
+        return self.bytes == other.bytes
+
+    def __add__(self, other):
+        return Size(self.bytes + other.bytes)
+
+    def __sub__(self, other):
+        return Size(self.bytes - other.bytes)
+
+    def __mul__(self, other):
+        return Size(self.bytes * int(other))
+
+    def __truediv__(self, other):
+        return Size(self.bytes / int(other))
+
+    def __floordiv__(self, other):
+        return Size(self.bytes // int(other))
+
+    def __rmul__(self, other):
+        return Size(self.bytes * int(other))
+
+    def __rtruediv__(self, other):
+        return Size(int(other) / self.bytes)
+
+    def __rfloordiv__(self, other):
+        return Size(int(other) // self.bytes)
+
+    def __iadd__(self, other):
+        self.bytes += other.bytes
+
+        return self
+
+    def __isub__(self, other):
+        self.bytes -= other.bytes
+
+        return self
+
+    def __imul__(self, other):
+        self.bytes *= int(other)
+
+        return self
+
+    def __itruediv__(self, other):
+        self.bytes /= int(other)
+
+        return self
+
+    def __ifloordir__(self, other):
+        self.bytes //= int(other)
+
+        return self
+
 
 def print_structure(struct, indent=0):
     print(struct)