Getting rid of the dependency in ruamel.yaml (#418)

zivnevo · web-flow · commit 1ad308ab150d · 2023-02-06T15:19:57.000+02:00
diff --git a/nca/CoreDS/Peer.py b/nca/CoreDS/Peer.py
@@ -128,7 +128,7 @@ def canonical_form(self):
 
         labels = sorted(self.labels.items())
         for label in labels:
-            ret += ',(' + label[0] + ',' + label[1] + ')'
+            ret += f',({label[0]},{label[1]})'
 
         return ret
 
diff --git a/nca/FileScanners/DirScanner.py b/nca/FileScanners/DirScanner.py
@@ -15,8 +15,8 @@ class DirScanner(GenericTreeScanner, HelmScanner):
        A class for reading yaml files from a file system path
     """
 
-    def __init__(self, fs_path, rt_load=False):
-        GenericTreeScanner.__init__(self, rt_load)
+    def __init__(self, fs_path, fast_load=False):
+        GenericTreeScanner.__init__(self, fast_load)
         HelmScanner.__init__(self)
         self.fs_path = fs_path
 
diff --git a/nca/FileScanners/GenericTreeScanner.py b/nca/FileScanners/GenericTreeScanner.py
@@ -5,9 +5,9 @@
 
 import os
 import abc
-from dataclasses import dataclass
-from ruamel.yaml import YAML, error
 from sys import stderr
+from dataclasses import dataclass
+import yaml
 
 
 @dataclass
@@ -19,15 +19,68 @@ class YamlFile:
     path: str
 
 
+class ObjectWithLocation:
+    line_number = 0
+    column_number = 0
+
+
+class YamlDict(dict, ObjectWithLocation):
+    pass
+
+
+class YamlList(list, ObjectWithLocation):
+    pass
+
+
+def to_yaml_objects(yaml_node):
+    if isinstance(yaml_node, yaml.SequenceNode):
+        res = YamlList()
+        res.line_number = yaml_node.start_mark.line
+        res.column_number = yaml_node.start_mark.column
+        for obj in yaml_node.value:
+            res.append(to_yaml_objects(obj))
+        return res
+    if isinstance(yaml_node, yaml.MappingNode):
+        res = YamlDict()
+        res.line_number = yaml_node.start_mark.line + 1
+        res.column_number = yaml_node.start_mark.column + 1
+        for obj in yaml_node.value:
+            res[obj[0].value] = to_yaml_objects(obj[1])
+        return res
+
+    # Node is a ScalarNode. First check if it can be interpreted as an int (e.g., port number)
+    try:
+        int_val = int(yaml_node.value)
+        return int_val
+    except ValueError:
+        pass
+
+    if not yaml_node.style:
+        # Now check if it is Boolean
+        if yaml_node.value == 'true':
+            return True
+        if yaml_node.value == 'false':
+            return False
+        # check if it's the null value
+        if yaml_node.value in ['', 'null']:
+            return None
+
+    return yaml_node.value
+
+
+def convert_documents(documents):
+    return [to_yaml_objects(document) for document in documents]
+
+
 class GenericTreeScanner(abc.ABC):
     """
     A base class for reading yaml files
     """
-    def __init__(self, rt_load=False):
+    def __init__(self, fast_load=False):
         """
-        :param bool rt_load: if True, load yaml with RoundTripLoader
+        :param bool fast_load: if True, load yaml faster, without saving objects location in file
         """
-        self.rt_load = rt_load
+        self.fast_load = fast_load
 
     @abc.abstractmethod
     def get_yamls(self):
@@ -49,14 +102,17 @@ def _yield_yaml_file(self, path, stream):
         :param str path: the path of the file
         :param stream: an IO-Text stream or Union of the file contents, depends on the scanner's type
         """
-        yaml = YAML(typ="rt") if self.rt_load else YAML(typ="safe")
         try:
-            yield YamlFile(yaml.load_all(stream), path)
-        except error.MarkedYAMLError as parse_error:
-            print(f'{parse_error.problem_mark.name}:{parse_error.problem_mark.line}:{parse_error.problem_mark.column}:',
+            if self.fast_load:
+                documents = yaml.load_all(stream, Loader=yaml.CSafeLoader)
+            else:
+                documents = convert_documents(yaml.compose_all(stream, Loader=yaml.CSafeLoader))
+            yield YamlFile(documents, path)
+        except yaml.MarkedYAMLError as parse_error:
+            print(f'{parse_error.problem_mark.name}:{parse_error.problem_mark.line+1}:{parse_error.problem_mark.column+1}:',
                   'Parse Error:', parse_error.problem, file=stderr)
             return
-        except error.YAMLError:
+        except yaml.YAMLError:
             print('Bad yaml file:', path, file=stderr)
             return
 
@@ -68,15 +124,15 @@ def _yield_yaml_file(self, path, stream):
 class TreeScannerFactory:
 
     @staticmethod
-    def get_scanner(entry, rt_load=False):
+    def get_scanner(entry, fast_load=False):
         """
         factory method to determine what scanner to build
         :param str entry: the entry (path/url) to be scanned
-        :param bool rt_load: if True, load yaml with RoundTripLoader
+        :param bool fast_load: if True, load yaml faster, without saving objects location in file
         """
         if entry.startswith(('https://github', GitScanner.raw_github_content_prefix)):
-            return GitScanner(entry, rt_load)
+            return GitScanner(entry, fast_load)
         elif os.path.isfile(entry) or os.path.isdir(entry) or (entry.endswith('**') and os.path.isdir(entry[:-2])):
-            return DirScanner(entry, rt_load)
+            return DirScanner(entry, fast_load)
         else:
             return None
diff --git a/nca/FileScanners/GitScanner.py b/nca/FileScanners/GitScanner.py
@@ -16,8 +16,8 @@ class GitScanner(GenericTreeScanner):
     """
     raw_github_content_prefix = 'https://raw.githubusercontent'
 
-    def __init__(self, url, rt_load=False):
-        GenericTreeScanner.__init__(self, rt_load)
+    def __init__(self, url, fast_load=False):
+        GenericTreeScanner.__init__(self, fast_load)
         self.url = url
         if url.startswith(self.raw_github_content_prefix):
             if not self.is_yaml_file(url):
diff --git a/nca/NetworkConfig/PoliciesFinder.py b/nca/NetworkConfig/PoliciesFinder.py
@@ -3,8 +3,8 @@
 # SPDX-License-Identifier: Apache2.0
 #
 
-from ruamel.yaml import YAML
 from collections import deque
+import yaml
 from nca.Utils.CmdlineRunner import CmdlineRunner
 from nca.Resources.NetworkPolicy import NetworkPolicy
 from nca.Parsers.K8sPolicyYamlParser import K8sPolicyYamlParser
@@ -122,8 +122,7 @@ def _add_policies_to_parse_queue(self, policy_list, file_name):
             self.parse_yaml_code_for_policy(policy, file_name)
 
     def _add_policies(self, doc, file_name):
-        yaml1 = YAML()
-        code = yaml1.load_all(doc)
+        code = yaml.load_all(doc, Loader=yaml.CSafeLoader)
         for policy_list in code:
             if isinstance(policy_list, dict):
                 self._add_policies_to_parse_queue(policy_list.get('items', []), file_name)
diff --git a/nca/NetworkConfig/ResourcesHandler.py b/nca/NetworkConfig/ResourcesHandler.py
@@ -5,15 +5,13 @@
 import copy
 import os
 from enum import Enum
-from sys import stderr
-from ruamel.yaml import error
 from nca.FileScanners.GenericTreeScanner import TreeScannerFactory
 from nca.Utils.CmdlineRunner import CmdlineRunner
+from nca.Utils.NcaLogger import NcaLogger
 from .NetworkConfig import NetworkConfig
 from .PoliciesFinder import PoliciesFinder
 from .TopologyObjectsFinder import PodsFinder, NamespacesFinder, ServicesFinder
 from .PeerContainer import PeerContainer
-from nca.Utils.NcaLogger import NcaLogger
 
 
 class ResourceType(Enum):
@@ -370,23 +368,15 @@ def parse_livesim_yamls(path):
             pods_finder = PodsFinder()
             ns_finder = NamespacesFinder()
             labels_found = {}
-            try:
-                for res_code in yaml_file.data:
-                    ns_finder.parse_yaml_code_for_ns(res_code)
-                    pods_finder.namespaces_finder = ns_finder
-                    pods_finder.add_eps_from_yaml(res_code)
-                for item in ns_finder.namespaces.values():
-                    labels_found.update(item.labels)
-                for item in pods_finder.peer_set:
-                    labels_found.update(item.labels)
-                results.update({yaml_file.path: labels_found})
-
-            except error.MarkedYAMLError as prs_err:
-                print(
-                    f'{prs_err.problem_mark.name}:{prs_err.problem_mark.line}:{prs_err.problem_mark.column}:',
-                    'Parse Error:', prs_err.problem, file=stderr)
-            except UnicodeDecodeError as decode_err:
-                print(f'Parse Error: Failed to decode {yaml_file.path}. error:\n{decode_err.reason}')
+            for res_code in yaml_file.data:
+                ns_finder.parse_yaml_code_for_ns(res_code)
+                pods_finder.namespaces_finder = ns_finder
+                pods_finder.add_eps_from_yaml(res_code)
+            for item in ns_finder.namespaces.values():
+                labels_found.update(item.labels)
+            for item in pods_finder.peer_set:
+                labels_found.update(item.labels)
+            results.update({yaml_file.path: labels_found})
 
         return results
 
@@ -406,32 +396,24 @@ def _parse_resources_path(self, resource_list, resource_flags):
             elif resource_item == 'istio':
                 self._handle_istio_inputs(resource_flags)
             else:
-                rt_load = True if ResourceType.Policies in resource_flags else False
-                resource_scanner = TreeScannerFactory.get_scanner(resource_item, rt_load=rt_load)
+                fast_load = ResourceType.Policies not in resource_flags
+                resource_scanner = TreeScannerFactory.get_scanner(resource_item, fast_load=fast_load)
                 if resource_scanner is None:
                     continue
                 yaml_files = resource_scanner.get_yamls()
                 if not yaml_files:
                     continue
                 for yaml_file in yaml_files:
-                    try:
-                        for res_code in yaml_file.data:
-                            if ResourceType.Namespaces in resource_flags:
-                                self.ns_finder.parse_yaml_code_for_ns(res_code)
-                            if ResourceType.Pods in resource_flags:
-                                self.pods_finder.namespaces_finder = self.ns_finder
-                                self.pods_finder.add_eps_from_yaml(res_code)
-                                self.services_finder.namespaces_finder = self.ns_finder
-                                self.services_finder.parse_yaml_code_for_service(res_code, yaml_file)
-                            if ResourceType.Policies in resource_flags:
-                                self.policies_finder.parse_yaml_code_for_policy(res_code, yaml_file.path)
-
-                    except error.MarkedYAMLError as prs_err:
-                        print(
-                            f'{prs_err.problem_mark.name}:{prs_err.problem_mark.line}:{prs_err.problem_mark.column}:',
-                            'Parse Error:', prs_err.problem, file=stderr)
-                    except UnicodeDecodeError as decode_err:
-                        print(f'Parse Error: Failed to decode {yaml_file.path}. error:\n{decode_err.reason}')
+                    for res_code in yaml_file.data:
+                        if ResourceType.Namespaces in resource_flags:
+                            self.ns_finder.parse_yaml_code_for_ns(res_code)
+                        if ResourceType.Pods in resource_flags:
+                            self.pods_finder.namespaces_finder = self.ns_finder
+                            self.pods_finder.add_eps_from_yaml(res_code)
+                            self.services_finder.namespaces_finder = self.ns_finder
+                            self.services_finder.parse_yaml_code_for_service(res_code, yaml_file)
+                        if ResourceType.Policies in resource_flags:
+                            self.policies_finder.parse_yaml_code_for_policy(res_code, yaml_file.path)
 
         self.policies_finder.parse_policies_in_parse_queue()
 
diff --git a/nca/NetworkConfig/TopologyObjectsFinder.py b/nca/NetworkConfig/TopologyObjectsFinder.py
@@ -31,11 +31,11 @@ def __init__(self):
     def load_peer_from_calico_resource(self):
         for peer_type in ['wep', 'hep', 'networkset', 'globalnetworkset']:
             peer_code = yaml.load(CmdlineRunner.get_calico_resources(peer_type),
-                                  Loader=yaml.SafeLoader)
+                                  Loader=yaml.CSafeLoader)
             self.add_eps_from_yaml(peer_code)
 
     def load_peer_from_k8s_live_cluster(self):
-        peer_code = yaml.load(CmdlineRunner.get_k8s_resources('pod'), Loader=yaml.SafeLoader)
+        peer_code = yaml.load(CmdlineRunner.get_k8s_resources('pod'), Loader=yaml.CSafeLoader)
         self.add_eps_from_yaml(peer_code)
 
     def add_eps_from_yaml(self, yaml_obj, kind_override=None):
@@ -262,7 +262,7 @@ def parse_yaml_code_for_ns(self, res_code):
 
     def load_ns_from_live_cluster(self):
         yaml_file = CmdlineRunner.get_k8s_resources('namespace')
-        ns_code = yaml.load(yaml_file, Loader=yaml.SafeLoader)
+        ns_code = yaml.load(yaml_file, Loader=yaml.CSafeLoader)
         self.set_namespaces(ns_code)
 
     def set_namespaces(self, ns_list):
@@ -326,7 +326,7 @@ def load_services_from_live_cluster(self):
         :return: The list of parsed services in K8sService format
         """
         yaml_file = CmdlineRunner.get_k8s_resources('service')
-        srv_resources = yaml.load(yaml_file, Loader=yaml.SafeLoader)
+        srv_resources = yaml.load(yaml_file, Loader=yaml.CSafeLoader)
         if not isinstance(srv_resources, dict):
             return
         parser = K8sServiceYamlParser('k8s')
diff --git a/nca/Parsers/CalicoPolicyYamlParser.py b/nca/Parsers/CalicoPolicyYamlParser.py
@@ -4,7 +4,6 @@
 #
 
 import re
-from ruamel.yaml import comments
 from nca.CoreDS.ProtocolNameResolver import ProtocolNameResolver
 from nca.CoreDS.Peer import PeerSet, IpBlock
 from nca.CoreDS.PortSet import PortSet
@@ -258,20 +257,6 @@ def _parse_port(self, port, array):
                 res_port_set.add_port(port)
         return res_port_set
 
-    @staticmethod
-    def _get_value_as_str(dict_to_use, key):
-        """
-        Safely getting string values that contain '!'.
-        Calico uses '!' as negation operator in selectors, while YAML uses it to declare tags
-        :param dict dict_to_use: The dictionary to retrieve the value from
-        :param str key: The key for which the value should be retrieved
-        :return: The proper string value
-        """
-        val = dict_to_use.get(key)
-        if isinstance(val, comments.TaggedScalar):  # negation operator '!' is used to declare tags in YAML
-            val = val.tag.value
-        return val
-
     def _get_rule_peers(self, entity_rule):
         """
         Parse the peer-specifying parts of the source/destination parts of a rule
@@ -291,9 +276,9 @@ def _get_rule_peers(self, entity_rule):
         for cidr in not_nets:
             rule_ips -= IpBlock(cidr)
 
-        ns_selector = self._get_value_as_str(entity_rule, 'namespaceSelector')
-        pod_selector = self._get_value_as_str(entity_rule, 'selector')
-        not_pod_selector = self._get_value_as_str(entity_rule, 'notSelector')
+        ns_selector = entity_rule.get('namespaceSelector')
+        pod_selector = entity_rule.get('selector')
+        not_pod_selector = entity_rule.get('notSelector')
         if ns_selector:
             rule_peers = self._parse_label_selector(ns_selector, entity_rule, namespace_selector=True)
         elif pod_selector:
diff --git a/nca/Parsers/GenericYamlParser.py b/nca/Parsers/GenericYamlParser.py
@@ -4,15 +4,15 @@
 #
 
 from sys import stderr
-from ruamel.yaml import comments
 from enum import Enum
 from nca.CoreDS.DimensionsManager import DimensionsManager
 from nca.CoreDS.TcpLikeProperties import TcpLikeProperties
 from nca.CoreDS.MethodSet import MethodSet
 from nca.CoreDS.ConnectionSet import ConnectionSet
 from nca.CoreDS.PortSet import PortSet
-from nca.Utils.NcaLogger import NcaLogger
 from nca.CoreDS.Peer import IpBlock
+from nca.Utils.NcaLogger import NcaLogger
+from nca.FileScanners.GenericTreeScanner import ObjectWithLocation
 
 
 class GenericYamlParser:
@@ -52,8 +52,8 @@ def syntax_error(self, msg, obj=None):
         :param obj: optionally, a CommentedBase object with context
         :return: None
         """
-        if isinstance(obj, comments.CommentedBase):
-            raise SyntaxError(msg, (self.yaml_file_name, obj.lc.line, obj.lc.col, '')) from None
+        if isinstance(obj, ObjectWithLocation):
+            raise SyntaxError(msg, (self.yaml_file_name, obj.line_number, obj.column_number, '')) from None
         raise SyntaxError(msg) from None
 
     def warning(self, msg, obj=None):
@@ -64,8 +64,8 @@ def warning(self, msg, obj=None):
         :return: None
         """
         print_msg = 'Warning: ' + msg
-        if isinstance(obj, comments.CommentedBase):
-            print_msg = f'{self.yaml_file_name}:{obj.lc.line}:{obj.lc.col}: {print_msg}'
+        if isinstance(obj, ObjectWithLocation):
+            print_msg = f'{self.yaml_file_name}:{obj.line_number}:{obj.column_number}: {print_msg}'
 
         NcaLogger().log_message(print_msg, file=stderr)
         self.warning_msgs.append(msg)
diff --git a/nca/SchemeRunner.py b/nca/SchemeRunner.py
diff --git a/requirements.txt b/requirements.txt
diff --git a/setup.cfg b/setup.cfg
diff --git a/tests/run_all_tests.py b/tests/run_all_tests.py
