Extract get_input_path function to be used by all scripts

tracetools_analysis/tracetools_analysis/scripts/__init__.py

@@ -0,0 +1,25 @@
+# Copyright 2019 Apex.AI, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+from typing import List
+
+
+def get_input_path(
+    argv: List[str] = sys.argv,
+) -> str:
+    if len(argv) < 2:
+        print('Syntax: [trace directory | converted tracefile]')
+        sys.exit(1)
+    return argv[1]

tracetools_analysis/tracetools_analysis/scripts/cb_durations.py

@@ -13,14 +13,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import sys
-
 import pandas as pd
 
 from tracetools_analysis.loading import load_file
 from tracetools_analysis.processor.ros2 import Ros2Handler
 from tracetools_analysis.utils.ros2 import Ros2DataModelUtil
 
+from . import get_input_path
+
 
 removals = [
     'void (', 'rclcpp::', 'std::shared_ptr<', '>', '::msg'
@@ -40,11 +40,9 @@ def format_fn(fname: str):
 
 
 def main():
-    if len(sys.argv) < 2:
-        print('Syntax: <tracefile>')
-        sys.exit(-1)
+    input_path = get_input_path()
 
-    events = load_file(sys.argv[1])
+    events = load_file(input_path)
     handler = Ros2Handler.process(events)
     du = Ros2DataModelUtil(handler.data)
 

tracetools_analysis/tracetools_analysis/scripts/memory_usage.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import sys
-
 from tracetools_analysis.loading import load_file
 from tracetools_analysis.processor import Processor
 from tracetools_analysis.processor.memory_usage import KernelMemoryUsageHandler
@@ -22,12 +20,11 @@ from tracetools_analysis.processor.ros2 import Ros2Handler
 from tracetools_analysis.utils.memory_usage import MemoryUsageDataModelUtil
 from tracetools_analysis.utils.ros2 import Ros2DataModelUtil
 
+from . import get_input_path
+
 
 def main():
-    if len(sys.argv) < 2:
-        print('Syntax: [trace directory | converted tracefile]')
-        sys.exit(1)
-    input_path = sys.argv[1]
+    input_path = get_input_path()
 
     events = load_file(input_path, do_convert_if_needed=True)
     ust_memory_handler = UserspaceMemoryUsageHandler()