Add to_directory method to relevant object classes

[oam.git] / code / network-generator / network_generation / view / network_viewer.py

diff --git a/code/network-generator/network_generation/view/network_viewer.py b/code/network-generator/network_generation/view/network_viewer.py
index 54ba84f..750a661 100644 (file)
--- a/code/network-generator/network_generation/view/network_viewer.py
+++ b/code/network-generator/network_generation/view/network_viewer.py
@@ -12,26 +12,29 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-#!/usr/bin/python
+# !/usr/bin/python
 """
 Provides functions to convert the Network into different formats
 """
 
+import gzip
 import json
+from typing_extensions import Buffer
+import zipfile
 import xml.etree.ElementTree as ET
+from typing import Any
 
 from network_generation.model.python.o_ran_network import ORanNetwork
 
 
 class NetworkViewer:
     """
-    This class contains all functions converting the Network into different formats
+    This class contains all functions converting the Network into
+    different formats
     """
 
-    __network: ORanNetwork = None
-
     # constructor
-    def __init__(self, network: ORanNetwork):
+    def __init__(self, network: ORanNetwork) -> None:
         self.__network = network
 
     # json format
@@ -43,7 +46,7 @@ class NetworkViewer:
         """
         return self
 
-    def show_as_json(self) -> dict[str, dict]:
+    def show_as_json(self) -> None:
         """
         Method printing the class in json format.
         """
@@ -55,16 +58,31 @@ class NetworkViewer:
         """
         print(self.__network)
 
-    def save(self, filename: str) -> None:
+    def to_directory(self, parent_dir: str) -> None:
+        """
+        Method converting the network to a subdirectory file structure
+        """
+        self.__network.to_directory(parent_dir)
+        print(f'Directory structure saved to "{parent_dir}"')
+
+    def save(self, filename: str, compressed: bool = True) -> None:
         """
         Method saving the class content to a file in json format.
         :param filename: A valid path to a file on the system.
+        :param compressed: if True, svg is stored as svgz format.
         :type filename: string
         """
-        with open(filename, "w", encoding="utf-8") as json_file:
-            output: dict[str, dict] = self.__network.to_topology()
-            json.dump(output, json_file, ensure_ascii=False, indent=2)
-            print("File '" + filename + "' saved!")
+        output: dict[str, Any] = self.__network.to_topology()
+        if compressed is True:
+            with gzip.open(f'{filename}-operational.json.gz', 'wb') as zf:
+                zf.write(json.dumps(output, indent=2).encode('utf-8'))
+            print(f'File "{filename}-operational.json.gz" saved!')
+        else:
+            with open(
+                f'{filename}-operational.json', "w", encoding="utf-8"
+            ) as jf:
+                json.dump(output, jf, ensure_ascii=False, indent=2)
+            print(f'File "{filename}-operational.json" saved!')
 
     def readStylesFromFile(self) -> str:
         """
@@ -75,27 +93,37 @@ class NetworkViewer:
             content = styles.read()
             return content
 
-    def svg(self, filename: str) -> None:
+    def svg(self, filename: str, compressed: bool = True) -> None:
         """
         Method saving the class content to a file in xml/svg format.
 
         :param filename: A valid path to a file on the system.
+        :param compressed: if True, svg is stored as svgz format.
         :type filename: string
         """
         root = self.__network.toSvg()
         style = ET.Element("style")
         style.text = self.readStylesFromFile()
         root.findall(".//desc")[0].append(style)
-        ET.ElementTree(root).write(
-            filename, encoding="utf-8", xml_declaration=True
-        )
-        print("File '" + filename + "' saved!")
 
-    def kml(self, filename: str) -> None:
+        if compressed is True:
+            svg_output: Buffer = ET.tostring(
+                root, encoding="utf-8", xml_declaration=True)
+            with gzip.open(f'{filename}.svgz', 'wb') as zf:
+                zf.write(svg_output)
+            print(f'File "{filename}.svgz" saved!')
+        else:
+            ET.ElementTree(root).write(
+                f'{filename}.svg', encoding="utf-8", xml_declaration=True
+            )
+            print(f'File "{filename}.svg" saved!')
+
+    def kml(self, filename: str, compressed: bool = True) -> None:
         """
         Method saving the class content to a file in xml/kml format.
 
         :param filename: A valid path to a file on the system.
+        :param compressed: if True, kml is stored as kmz format.
         :type filename: string
         """
         root = self.__network.toKml()
@@ -106,15 +134,24 @@ class NetworkViewer:
                 style = ET.Element("Style", {"id": key})
                 line_style = ET.SubElement(style, "LineStyle")
                 color = ET.SubElement(line_style, "color")
-                color.text = value["stroke"]["color"]
+                color.text = str(value["stroke"]["color"])
                 width = ET.SubElement(line_style, "width")
-                width.text = value["stroke"]["width"]
+                width.text = str(value["stroke"]["width"])
                 poly_style = ET.SubElement(style, "PolyStyle")
                 fill = ET.SubElement(poly_style, "color")
-                fill.text = value["fill"]["color"]
+                fill.text = str(value["fill"]["color"])
                 root.findall(".//Document")[0].append(style)
 
-        ET.ElementTree(root).write(
-            filename, encoding="utf-8", xml_declaration=True
-        )
-        print("File '" + filename + "' saved!")
+        kml: str = ET.tostring(
+            root, encoding="utf-8", xml_declaration=True)
+        if compressed is True:
+            with zipfile.ZipFile(
+                f'{filename}.kmz', 'w', zipfile.ZIP_DEFLATED
+            ) as zf:
+                zf.writestr(f'{filename.split("/")[1]}.kml', data=kml)
+            print(f'File "{filename}.kmz" saved!')
+        else:
+            kml_file = open(f'{filename}.kml', 'w')
+            kml_file.write(kml)
+            kml_file.close()
+            print(f'File "{filename}.kml" saved!')