Merge pull request #459 from piotrkochan/feature/cap-parse-ast-json

capabilities_parser parse ast and json

Author: mdmintz

seleniumbase/core/capabilities_parser.py

@@ -1,18 +1,35 @@
 import re
+import ast
+import json
 
 
-def get_desired_capabilities(cap_file):
-    if not cap_file.endswith('.py'):
-        raise Exception("\n\n`%s` is not a Python file!\n\n" % cap_file)
+def _analyze_ast(contents):
+    try:
+        return ast.literal_eval(contents)
+    except SyntaxError:
+        pass
+    try:
+        # remove all comments
+        contents = re.sub(re.compile(r"/\*.*?\*/", re.DOTALL), "", contents)
+        contents = re.sub(re.compile(r"#.*?\n"), "", contents)
+
+        # remove anything before dict declaration like: "caps = { ..."
+        match = re.match(r"^([^{]+)", contents)
+        if match:
+            contents = contents.replace(match.group(1), "")
+
+        # and try again
+        return ast.literal_eval(contents)
+    except SyntaxError:
+        pass
+
+    return False
 
-    f = open(cap_file, 'r')
-    all_code = f.read()
-    f.close()
 
-    desired_capabilities = {}
-    num_capabilities = 0
+def _analyze_manual(contents):
+    capabilities = {}
 
-    code_lines = all_code.split('\n')
+    code_lines = contents.split('\n')
     for line in code_lines:
         if "desired_cap = {" in line:
             line = line.split("desired_cap = {")[1]
@@ -22,17 +39,15 @@ def get_desired_capabilities(cap_file):
         if data:
             key = data.group(1)
             value = data.group(2)
-            desired_capabilities[key] = value
-            num_capabilities += 1
+            capabilities[key] = value
             continue
 
         # "KEY" : "VALUE"
         data = re.match(r'^\s*"([\S\s]+)"\s*:\s*"([\S\s]+)"\s*[,}]?\s*$', line)
         if data:
             key = data.group(1)
             value = data.group(2)
-            desired_capabilities[key] = value
-            num_capabilities += 1
+            capabilities[key] = value
             continue
 
         # 'KEY' : "VALUE"
@@ -41,8 +56,7 @@ def get_desired_capabilities(cap_file):
         if data:
             key = data.group(1)
             value = data.group(2)
-            desired_capabilities[key] = value
-            num_capabilities += 1
+            capabilities[key] = value
             continue
 
         # "KEY" : 'VALUE'
@@ -51,8 +65,7 @@ def get_desired_capabilities(cap_file):
         if data:
             key = data.group(1)
             value = data.group(2)
-            desired_capabilities[key] = value
-            num_capabilities += 1
+            capabilities[key] = value
             continue
 
         # "KEY" : True
@@ -61,8 +74,7 @@ def get_desired_capabilities(cap_file):
         if data:
             key = data.group(1)
             value = True
-            desired_capabilities[key] = value
-            num_capabilities += 1
+            capabilities[key] = value
             continue
 
         # 'KEY' : True
@@ -71,8 +83,7 @@ def get_desired_capabilities(cap_file):
         if data:
             key = data.group(1)
             value = True
-            desired_capabilities[key] = value
-            num_capabilities += 1
+            capabilities[key] = value
             continue
 
         # "KEY" : False
@@ -81,8 +92,7 @@ def get_desired_capabilities(cap_file):
         if data:
             key = data.group(1)
             value = False
-            desired_capabilities[key] = value
-            num_capabilities += 1
+            capabilities[key] = value
             continue
 
         # 'KEY' : False
@@ -91,26 +101,23 @@ def get_desired_capabilities(cap_file):
         if data:
             key = data.group(1)
             value = False
-            desired_capabilities[key] = value
-            num_capabilities += 1
+            capabilities[key] = value
             continue
 
         # caps['KEY'] = 'VALUE'
         data = re.match(r"^\s*caps\['([\S\s]+)'\]\s*=\s*'([\S\s]+)'\s*$", line)
         if data:
             key = data.group(1)
             value = data.group(2)
-            desired_capabilities[key] = value
-            num_capabilities += 1
+            capabilities[key] = value
             continue
 
         # caps["KEY"] = "VALUE"
         data = re.match(r'^\s*caps\["([\S\s]+)"\]\s*=\s*"([\S\s]+)"\s*$', line)
         if data:
             key = data.group(1)
             value = data.group(2)
-            desired_capabilities[key] = value
-            num_capabilities += 1
+            capabilities[key] = value
             continue
 
         # caps['KEY'] = "VALUE"
@@ -119,8 +126,7 @@ def get_desired_capabilities(cap_file):
         if data:
             key = data.group(1)
             value = data.group(2)
-            desired_capabilities[key] = value
-            num_capabilities += 1
+            capabilities[key] = value
             continue
 
         # caps["KEY"] = 'VALUE'
@@ -129,8 +135,7 @@ def get_desired_capabilities(cap_file):
         if data:
             key = data.group(1)
             value = data.group(2)
-            desired_capabilities[key] = value
-            num_capabilities += 1
+            capabilities[key] = value
             continue
 
         # caps["KEY"] = True
@@ -139,8 +144,7 @@ def get_desired_capabilities(cap_file):
         if data:
             key = data.group(1)
             value = True
-            desired_capabilities[key] = value
-            num_capabilities += 1
+            capabilities[key] = value
             continue
 
         # caps['KEY'] = True
@@ -149,8 +153,7 @@ def get_desired_capabilities(cap_file):
         if data:
             key = data.group(1)
             value = True
-            desired_capabilities[key] = value
-            num_capabilities += 1
+            capabilities[key] = value
             continue
 
         # caps["KEY"] = False
@@ -159,8 +162,7 @@ def get_desired_capabilities(cap_file):
         if data:
             key = data.group(1)
             value = False
-            desired_capabilities[key] = value
-            num_capabilities += 1
+            capabilities[key] = value
             continue
 
         # caps['KEY'] = False
@@ -169,11 +171,45 @@ def get_desired_capabilities(cap_file):
         if data:
             key = data.group(1)
             value = False
-            desired_capabilities[key] = value
-            num_capabilities += 1
+            capabilities[key] = value
             continue
 
-    if num_capabilities == 0:
+    return capabilities
+
+
+def _read_file(file):
+    f = open(file, 'r')
+    data = f.read()
+    f.close()
+
+    return data
+
+
+def _parse_py_file(cap_file):
+    all_code = _read_file(cap_file)
+    capabilities = _analyze_ast(all_code)
+
+    if not capabilities:
+        capabilities = _analyze_manual(all_code)
+
+    return capabilities
+
+
+def _parse_json_file(cap_file):
+    all_code = _read_file(cap_file)
+
+    return json.loads(all_code)
+
+
+def get_desired_capabilities(cap_file):
+    if cap_file.endswith('.py'):
+        capabilities = _parse_py_file(cap_file)
+    elif cap_file.endswith('.json'):
+        capabilities = _parse_json_file(cap_file)
+    else:
+        raise Exception("\n\n`%s` is not a Python or JSON file!\n" % cap_file)
+
+    if len(capabilities.keys()) == 0:
         raise Exception("Unable to parse desired capabilities file!")
 
-    return desired_capabilities
+    return capabilities