Ejemplo n.º 1
0
 def init_save_database(self):
     from topic_store.database import MongoStorage
     self.db_client = MongoStorage(config=self.scenario.storage["config"],
                                   collection=self.scenario.context,
                                   use_grid_fs=self.use_grid_fs)
     self.info_logger("Initialised saving to database {} @ '{}/{}'".format(
         self.db_client.uri, self.db_client.name, self.scenario.context))
Ejemplo n.º 2
0
def mongodb_to_ros_bag(scenario_file, output_file):
    scenario = ScenarioFileParser(scenario_file)
    if scenario.storage["method"] != "database":
        raise ValueError(
            "Scenario file '{}' storage.method is not set to 'database'".
            format(scenario_file))
    print("Converting MongoDB '{}' to ROS bag '{}'".format(
        scenario.storage["uri"], output_file.name))
    client = MongoStorage(uri=scenario.storage["uri"],
                          collection=scenario.context)

    storage = get_mongo_storage_by_session(client)
    count = storage.cursor.count()

    ros_bag = rosbag.Bag(str(output_file), 'w')

    try:
        with tqdm(total=count) as progress_bar:
            for item in storage:
                msgs = item.to_ros_msg_list()
                time = rospy.Time.from_sec(item["_ts_meta"]["ros_time"])
                for msg in msgs:
                    if hasattr(msg, "_connection_header"):
                        source = getattr(msg, "_connection_header")["topic"]
                        if source:
                            try:
                                ros_bag.write(source, msg, time)
                            except Exception as e:
                                print("Could not write", source, 'because',
                                      e.message)
                progress_bar.update()
    finally:
        print("Closing the ROS bag '{}'".format(output_file))
        ros_bag.close()
Ejemplo n.º 3
0
    def test_image_encoding(self):
        if not rospy.get_node_uri():
            rospy.init_node("topic_store_tests")

        client = MongoStorage(collection="python_tests")

        # Test ROS_Message Serialisation with common Image Types
        from topic_store.utils import _numpy_types as numpy_types

        for encoding, np_type in numpy_types.items():
            size = (32, 32, np_type[1]) if np_type[1] != 1 else (32, 32)
            random_array = np.random.random(size)
            typed_array = (random_array * 255.0).astype(np_type[0])
            message = ros_numpy.msgify(Image, typed_array, encoding=encoding)

            # Insert image message
            im_document = TopicStore({"image": message})
            im_result = client.insert_one(im_document)

            # Get image message and check data is the same
            returned_im_document = client.find_by_id(im_result.inserted_id)
            assert returned_im_document.id == im_result.inserted_id
            retrieved_array = ros_numpy.numpify(
                returned_im_document.msgs["image"])
            np.testing.assert_equal(typed_array, retrieved_array)

            # Delete image message
            client.delete_by_id(im_result.inserted_id)
Ejemplo n.º 4
0
def mongodb_to_topic_store(scenario_file, topic_store_file):
    client = MongoStorage.load(scenario_file)
    print("Converting MongoDB '{}' to '{}'".format(client.uri, topic_store_file.name))

    storage = get_mongo_storage_by_session(client)
    count = storage.cursor.count()

    topic_storage = TopicStorage(topic_store_file)

    with tqdm(total=count) as progress_bar:
        for item in storage:
            topic_storage.insert_one(item)
            progress_bar.update()
Ejemplo n.º 5
0
def topic_store_to_mongodb(topic_store_file, scenario_file):
    client = MongoStorage.load(scenario_file)
    print("Converting '{}' to MongoDB '{}'".format(topic_store_file.name, client.uri))

    storage = TopicStorage.load(topic_store_file)
    count = len(storage)  # TODO: very slow operation
    with tqdm(total=count) as progress_bar:
        for item in storage:
            try:
                client.insert_one(item)
            except pymongo.errors.DuplicateKeyError:
                print("Storage Item '_id: {}' already exists in the '{}/{}' collection".format(item.id, client.name,
                                                                                               client.collection_name))
            progress_bar.update()
Ejemplo n.º 6
0
def client_from_uri(uri, collection):
    if is_uri(uri):
        # DB name will usually be specified as authSource in the URI, if not present use default=topic_store
        db_name = None
        if "authSource" in uri:
            options = [
                s.split('=') for s in urlparse(uri).query.split("&") if s
            ]
            options = {k: v for k, v in options}
            if "authSource" in options:
                db_name = options["authSource"]
        client = MongoStorage(collection=collection, uri=uri, db_name=db_name)
        return client
    else:
        raise ValueError("Not a valid URI: {}".format(uri))
Ejemplo n.º 7
0
    def test_documents(self):
        if not rospy.get_node_uri():
            rospy.init_node("topic_store_tests")

        client = MongoStorage(collection="python_tests")

        # Insert a test document
        insert_result = client.insert_one(
            TopicStore({
                "name": "test_name",
                "number": 1
            }))

        # Retrieve the inserted document
        inserted_document = client.find_by_id(insert_result.inserted_id)
        assert inserted_document.id == insert_result.inserted_id

        # Update the document name and number fields
        new_name = ''.join(random.sample('raymond', 7))
        new_number = random.randint(0, 100)
        update_result = client.update_one_by_id(inserted_document.id,
                                                name=new_name,
                                                number=new_number)

        inserted_document_after_update = client.find_by_id(
            insert_result.inserted_id)
        assert inserted_document_after_update.id == insert_result.inserted_id
        assert inserted_document_after_update.dict["number"] == new_number
        assert inserted_document_after_update.dict["name"] == new_name

        # Print all documents in the collection
        cursor = client.find()
        for x in cursor:
            print(
                "Doc:\n\t-As Structure: {}\n\t-As Dict: {}\n\t-As ROS Msgs: {}"
                .format(str(x), x.dict, x.msgs))

        # Or print using the same API as TopicStorage
        for x in client:
            print(
                "Doc:\n\t-As Structure: {}\n\t-As Dict: {}\n\t-As ROS Msgs: {}"
                .format(str(x), x.dict, x.msgs))

        # Cleanup test by deleting document
        delete_result = client.delete_by_id(insert_result.inserted_id)
Ejemplo n.º 8
0
    def test_large_document(self):
        if not rospy.get_node_uri():
            rospy.init_node("topic_store_tests")

        client = MongoStorage(collection="python_tests", use_grid_fs=True)

        # Insert >16MB document
        random_array = np.random.random((3000, 3000, 3)).astype(np.float32)
        message = ros_numpy.msgify(Image, random_array, encoding="32FC3")

        # Insert image message
        im_document = TopicStore({"image": message})
        im_result = client.insert_one(im_document)

        # Get image message and check data is the same
        returned_im_document = client.find_by_id(im_result.inserted_id)
        assert returned_im_document.id == im_result.inserted_id
        retrieved_array = ros_numpy.numpify(returned_im_document.msgs["image"])
        np.testing.assert_equal(random_array, retrieved_array)

        # Delete image message
        client.delete_by_id(im_result.inserted_id)
Ejemplo n.º 9
0
    def test_image_encoding(self):
        if not rospy.get_node_uri():
            rospy.init_node("topic_store_tests")

        client = MongoStorage(collection="python_tests")

        # Test ROS_Message Serialisation with common Image Types
        numpy_types = {
            "rgb8": (np.uint8, 3),
            "rgba8": (np.uint8, 4),
            "rgb16": (np.uint16, 3),
            "rgba16": (np.uint16, 4),
            "bgr8": (np.uint8, 3),
            "bgra8": (np.uint8, 4),
            "bgr16": (np.uint16, 3),
            "bgra16": (np.uint16, 4),
            "mono8": (np.uint8, 1),
            "mono16": (np.uint16, 1),
            "bayer_rggb8": (np.uint8, 1),
            "bayer_bggr8": (np.uint8, 1),
            "bayer_gbrg8": (np.uint8, 1),
            "bayer_grbg8": (np.uint8, 1),
            "bayer_rggb16": (np.uint16, 1),
            "bayer_bggr16": (np.uint16, 1),
            "bayer_gbrg16": (np.uint16, 1),
            "bayer_grbg16": (np.uint16, 1),
            "8UC1": (np.uint8, 1),
            "8UC2": (np.uint8, 2),
            "8UC3": (np.uint8, 3),
            "8UC4": (np.uint8, 4),
            "8SC1": (np.int8, 1),
            "8SC2": (np.int8, 2),
            "8SC3": (np.int8, 3),
            "8SC4": (np.int8, 4),
            "16UC1": (np.uint16, 1),
            "16UC2": (np.uint16, 2),
            "16UC3": (np.uint16, 3),
            "16UC4": (np.uint16, 4),
            "16SC1": (np.int16, 1),
            "16SC2": (np.int16, 2),
            "16SC3": (np.int16, 3),
            "16SC4": (np.int16, 4),
            "32SC1": (np.int32, 1),
            "32SC2": (np.int32, 2),
            "32SC3": (np.int32, 3),
            "32SC4": (np.int32, 4),
            "32FC1": (np.float32, 1),
            "32FC2": (np.float32, 2),
            "32FC3": (np.float32, 3),
            "32FC4": (np.float32, 4),
            "64FC1": (np.float64, 1),
            "64FC2": (np.float64, 2),
            "64FC3": (np.float64, 3),
            "64FC4": (np.float64, 4)
        }

        for encoding, np_type in numpy_types.items():
            size = (32, 32, np_type[1]) if np_type[1] != 1 else (32, 32)
            random_array = np.random.random(size)
            typed_array = (random_array * 255.0).astype(np_type[0])
            message = ros_numpy.msgify(Image, typed_array, encoding=encoding)

            # Insert image message
            im_document = TopicStore({"image": message})
            im_result = client.insert_one(im_document)

            # Get image message and check data is the same
            returned_im_document = client.find_by_id(im_result.inserted_id)
            assert returned_im_document.id == im_result.inserted_id
            retrieved_array = ros_numpy.numpify(
                returned_im_document.msgs["image"])
            np.testing.assert_equal(typed_array, retrieved_array)

            # Delete image message
            client.delete_by_id(im_result.inserted_id)
Ejemplo n.º 10
0
class ScenarioRunner:
    def __init__(self,
                 scenario_file,
                 stabilise_time,
                 verbose=True,
                 queue_size=100,
                 n_io_threads=1,
                 threads_auto=False,
                 use_grid_fs=False):
        self.saved_n = 0

        self.scenario_file = scenario_file
        self.stabilise_time = stabilise_time
        self.verbose = verbose
        self.thread_queue_size = queue_size
        self.n_threads = n_io_threads
        self.auto_threads = threads_auto
        self.use_grid_fs = use_grid_fs
        self.events = {}

        # Load Scenario
        rospy.loginfo("Loading scenario: '{}'".format(scenario_file))
        self.scenario = ScenarioFileParser(scenario_file)

        # Create publisher for topic_store scenario logs
        self.logger = best_logger(verbose=verbose)
        self.info_logger = DefaultLogger(verbose=verbose)

        # Create subscriber tree for getting all the data
        self.subscriber_tree = SubscriberTree(self.scenario.data)
        if self.stabilise_time:
            self.info_logger(
                "Waiting for {}s before starting (stabilise_time)".format(
                    self.stabilise_time))
            rospy.sleep(self.stabilise_time)

        # Choose appropriate methods
        # Setting up the scenario runner this way means it's easily extendable by inheriting from Scenario runner
        # and declaring custom behaviours
        save_init_function_name = "init_save_" + self.scenario.storage["method"]
        self.save_method_init_function = getattr(self, save_init_function_name,
                                                 None)
        if not callable(self.save_method_init_function):
            raise Exception(
                "Invalid storage value ('{}()' does not exist)".format(
                    save_init_function_name))
        self.save_method_init_function()

        save_function_name = "save_" + self.scenario.storage["method"]
        self.save_method_function = getattr(self, save_function_name, None)
        if not callable(self.save_method_function):
            raise Exception(
                "Invalid storage value ('{}()' does not exist)".format(
                    save_function_name))

        collection_method_init_name = "init_way_point_" + self.scenario.collection[
            "method"]
        self.collection_method_init_function = getattr(
            self, collection_method_init_name, None)
        if not callable(self.collection_method_init_function):
            raise Exception(
                "Invalid way point value ('{}()' does not exist)".format(
                    collection_method_init_name))

        self.save_callback_rate = FPSCounter()
        self.event_callback_rate = FPSCounter()
        self.extra_logs = {}
        if self.n_threads > 0:
            self.jobs_worker = LoadBalancer(maxsize=self.thread_queue_size,
                                            threads=self.n_threads,
                                            auto=self.auto_threads)
            self.save_callback_rate = FPSCounter(self.jobs_worker.maxsize)
        self.collection_method_init_function()

    def init_way_point_action_server(self):
        def __request(goal_msg):
            success, save_msg = True, "Success!"
            # TODO: self.save() no longer returns values of self.__save() because jobs are handled in a load_balancer
            # TODO: for now action_server will return success no matter what
            self.save()
            result = CollectDataResult(success)
            feedback = CollectDataFeedback(save_msg)
            self.service_server.publish_feedback(feedback)
            (self.service_server.set_succeeded
             if success else self.service_server.set_aborted)(result)

        action_lib_server_name = self.scenario.collection["action_server_name"]
        self.info_logger(
            "Starting '{}' actionlib server".format(action_lib_server_name))
        self.service_server = actionlib.SimpleActionServer(
            action_lib_server_name, CollectDataAction, __request, False)
        self.service_server.start()

    def init_way_point_action_server_video(self):
        event_name = "action_lib_start_stop"
        topic_to_watch = self.scenario.collection["watch_topic"]

        self.events[event_name] = {"event": Event(), "data": False}
        self.events[topic_to_watch] = {"event": Event(), "data": ""}

        def __request(goal_msg):
            should_start = str(
                goal_msg.message).lower() in ["true", "t", "start"]
            starting_string = "Starting" if should_start else "Stopping"
            self.info_logger("{} data collection from '{}'".format(
                starting_string, action_lib_server_name))
            self.set_event_msg_callback(should_start, event_name)

            result = CollectDataResult(should_start)
            feedback = CollectDataFeedback(
                "{} data capture".format(starting_string))
            self.service_server.publish_feedback(feedback)
            self.service_server.set_succeeded(result)

        action_lib_server_name = self.scenario.collection["action_server_name"]
        self.info_logger(
            "Starting '{}' actionlib server".format(action_lib_server_name))
        self.service_server = actionlib.SimpleActionServer(
            action_lib_server_name, CollectDataAction, __request, False)
        self.service_server.start()

        AutoSubscriber(topic_to_watch,
                       callback=self.set_event_msg_callback,
                       callback_args=topic_to_watch)

        while not rospy.is_shutdown():
            self.events[topic_to_watch]["event"].wait()
            self.events[topic_to_watch]["event"].clear()
            if self.events[event_name]["data"]:
                self.save()
                # self.log("Waiting for event on '{}' topic before next data cycle".format(topic_to_watch))

    def set_event_msg_callback(self, data, event_id):
        if event_id in self.events:
            self.event_callback_rate.toc()
            self.event_callback_rate.tic()
            self.extra_logs["event_callback_rate"] = "{:.2f}".format(
                self.save_callback_rate.get_fps())
            self.events[event_id]["data"] = data
            self.events[event_id]["event"].set()

    def init_way_point_timer(self):
        delay = self.scenario.collection["timer_delay"]
        while not rospy.is_shutdown():
            self.save()
            # self.log("Waiting for {}s before next data cycle".format(delay))
            rospy.sleep(delay)

    def init_way_point_event(self):
        event_name = "topic_update"
        topic_to_watch = self.scenario.collection["watch_topic"]
        self.events[event_name] = {"event": Event(), "data": ""}
        AutoSubscriber(topic_to_watch,
                       callback=self.set_event_msg_callback,
                       callback_args=event_name)
        while not rospy.is_shutdown():
            self.events[event_name]["event"].wait()
            self.events[event_name]["event"].clear()
            self.save()
            # self.log("Waiting for event on '{}' topic before next data cycle".format(topic_to_watch))

    def init_save_database(self):
        from topic_store.database import MongoStorage
        self.db_client = MongoStorage(config=self.scenario.storage["config"],
                                      collection=self.scenario.context,
                                      use_grid_fs=self.use_grid_fs)
        self.info_logger("Initialised saving to database {} @ '{}/{}'".format(
            self.db_client.uri, self.db_client.name, self.scenario.context))

    def init_save_filesystem(self):
        from topic_store.filesystem import TopicStorage
        formatted_datetime = datetime.datetime.now().strftime(
            '%Y_%m_%d_%H_%M_%S')
        save_location = self.scenario.storage["location"]
        if not save_location or save_location in [
                "default", "auto", "topic_store"
        ]:
            save_location = pathlib.Path(
                os.path.expanduser("~/.ros/topic_store/filesystem"))
        elif save_location.startswith("pkg="):
            package_name = save_location.split('=')[-1]
            save_location = pathlib.Path(
                os.path.expanduser(
                    "~/.ros/topic_store/filesystem/{}/".format(package_name)))
        else:
            save_location = pathlib.Path(os.path.expanduser(save_location))
        save_folder = save_location / self.scenario.context
        self.info_logger(
            "Configured save_location as '{}'".format(save_folder))
        try:
            save_folder.mkdir(parents=True)
        except OSError as e:
            if e.errno != 17:  # File exists is okay
                raise

        save_file = save_folder / "{}{}".format(formatted_datetime,
                                                TopicStorage.suffix)
        self.filesystem_storage = TopicStorage(save_file)
        self.info_logger("Initialised saving to the filesystem at '{}'".format(
            self.filesystem_storage.path))

    def save_database(self, message_tree):
        insert_result = self.db_client.insert_one(message_tree)
        self.saved_n += 1
        # self.log("Inserted document to database result='acknowledged={}, inserted_id={}'".format(
        #     insert_result.acknowledged, insert_result.inserted_id))
        return insert_result.inserted_id

    def save_filesystem(self, message_tree):
        # self.log("Saving documents to file system n={}".format(self.saved_n))
        self.saved_n += 1
        self.filesystem_storage.insert_one(message_tree)
        return self.saved_n

    def __save(self, data, job_meta=None, **kwargs):
        try:
            saved_data_id = self.save_method_function(data)
        except Exception as e:
            self.info_logger("Exception raised when saving! '{}'".format(
                e.message))
            return False, e.message
        if job_meta is not None:
            worker_id = job_meta.pop("worker_id", None)
            # Pretty print some variables
            info_str = {
                k.replace("worker_", ""): "{:.2f}".format(v)
                for k, v in job_meta.items()
            }
            self.extra_logs["save_callback_rate"] = "{:.2f}".format(
                self.save_callback_rate.get_fps())
            for k, v in self.extra_logs.items():
                info_str[k] = v
            self.logger(
                "Worker {} successfully saved data id='{}'".format(
                    worker_id, saved_data_id), **info_str)
        else:
            self.logger(
                "Successfully saved data id='{}'".format(saved_data_id),
                **self.extra_logs)
        return True, "Success!"

    def save(self):
        """Collates data from the scenario topic structure and saves. Returns SaveSuccess, SaveMessage"""
        self.save_callback_rate.toc()
        self.save_callback_rate.tic()
        data_hierarchy = self.subscriber_tree.get_message_tree()
        if self.n_threads == 0:  # don't use a threading model
            self.__save(data_hierarchy)
        else:
            added_task = self.jobs_worker.add_task(self.__save,
                                                   [data_hierarchy], {},
                                                   wait=False)
            if not added_task:
                self.info_logger(
                    "IO Queue Full, cannot add jobs to the queue, please wait.",
                    verbose=True)
                self.jobs_worker.add_task(self.__save, [data_hierarchy], {},
                                          wait=True)
Ejemplo n.º 11
0
class ScenarioRunner:
    def __init__(self, scenario_file, stabilise_time, verbose=True):
        self.saved_n = 0

        self.scenario_file = scenario_file
        self.stabilise_time = stabilise_time
        self.verbose = verbose
        self.logger = print
        self.events = {}

        # Load Scenario
        rospy.loginfo("Loading scenario: '{}'".format(scenario_file))
        self.scenario = ScenarioFileParser(scenario_file)

        # Create subscriber tree for getting all the data
        self.subscriber_tree = SubscriberTree(self.scenario.data)
        if self.stabilise_time:
            self.log("Waiting for {}s before starting (stabilise_time)".format(self.stabilise_time))
            rospy.sleep(self.stabilise_time)

        # Choose appropriate methods
        # Setting up the scenario runner this way means it's easily extendable by inheriting from Scenario runner
        # and declaring custom behaviours
        save_init_function_name = "init_save_" + self.scenario.storage["method"]
        self.save_method_init_function = getattr(self, save_init_function_name, None)
        if not callable(self.save_method_init_function):
            raise Exception("Invalid storage value ('{}()' does not exist)".format(save_init_function_name))
        self.save_method_init_function()

        save_function_name = "save_" + self.scenario.storage["method"]
        self.save_method_function = getattr(self, save_function_name, None)
        if not callable(self.save_method_function):
            raise Exception("Invalid storage value ('{}()' does not exist)".format(save_function_name))

        collection_method_init_name = "init_way_point_" + self.scenario.collection["method"]
        self.collection_method_init_function = getattr(self, collection_method_init_name, None)
        if not callable(self.collection_method_init_function):
            raise Exception("Invalid way point value ('{}()' does not exist)".format(collection_method_init_name))

        self.collection_method_init_function()

    def log(self, message, **kwargs):
        if self.verbose:
            self.logger("\033[93mScenarioRunner\033[0m: {}".format(message), **kwargs)

    def init_way_point_action_server(self):
        def __request(goal_msg):
            success, save_msg = self.save()
            result = CollectDataResult(success)
            feedback = CollectDataFeedback(save_msg)
            self.service_server.publish_feedback(feedback)
            (self.service_server.set_succeeded if success else self.service_server.set_aborted)(result)

        action_lib_server_name = self.scenario.collection["action_server_name"]
        self.log("Starting '{}' actionlib server".format(action_lib_server_name))
        self.service_server = actionlib.SimpleActionServer(action_lib_server_name, CollectDataAction, __request, False)
        self.service_server.start()

    def init_way_point_action_server_video(self):
        event_name = "action_lib_start_stop"
        topic_to_watch = self.scenario.collection["watch_topic"]

        self.events[event_name] = {"event": Event(), "data": False}
        self.events[topic_to_watch] = {"event": Event(), "data": ""}

        def __request(goal_msg):
            should_start = str(goal_msg.message).lower() in ["true", "t", "start"]
            starting_string = "Starting" if should_start else "Stopping"
            self.log("{} data collection from '{}'".format(starting_string, action_lib_server_name))
            self.set_event_msg_callback(should_start, event_name)

            result = CollectDataResult(should_start)
            feedback = CollectDataFeedback("{} data capture".format(starting_string))
            self.service_server.publish_feedback(feedback)
            self.service_server.set_succeeded(result)

        action_lib_server_name = self.scenario.collection["action_server_name"]
        self.log("Starting '{}' actionlib server".format(action_lib_server_name))
        self.service_server = actionlib.SimpleActionServer(action_lib_server_name, CollectDataAction, __request, False)
        self.service_server.start()

        AutoSubscriber(topic_to_watch, callback=self.set_event_msg_callback, callback_args=topic_to_watch)

        while not rospy.is_shutdown():
            self.events[topic_to_watch]["event"].wait()
            self.events[topic_to_watch]["event"].clear()
            if self.events[event_name]["data"]:
                self.save()
                self.log("Waiting for event on '{}' topic before next data cycle".format(topic_to_watch))

    def set_event_msg_callback(self, data, event_id):
        if event_id in self.events:
            self.events[event_id]["data"] = data
            self.events[event_id]["event"].set()

    def init_way_point_timer(self):
        delay = self.scenario.collection["timer_delay"]
        while not rospy.is_shutdown():
            self.save()
            self.log("Waiting for {}s before next data cycle".format(delay))
            rospy.sleep(delay)

    def init_way_point_event(self):
        event_name = "topic_update"
        topic_to_watch = self.scenario.collection["watch_topic"]
        self.events[event_name] = {"event": Event(), "data": ""}
        AutoSubscriber(topic_to_watch, callback=self.set_event_msg_callback, callback_args=event_name)
        while not rospy.is_shutdown():
            self.events[event_name]["event"].wait()
            self.events[event_name]["event"].clear()
            self.save()
            self.log("Waiting for event on '{}' topic before next data cycle".format(topic_to_watch))

    def init_save_database(self):
        from topic_store.database import MongoStorage
        self.db_client = MongoStorage(config=self.scenario.storage["config"], collection=self.scenario.context)
        self.log("Initialised saving to database {} @ '{}/{}'".format(self.db_client.uri,
                                                                      self.db_client.name, self.scenario.context))

    def init_save_filesystem(self):
        from topic_store.filesystem import TopicStorage
        formatted_datetime = datetime.datetime.now().strftime('%Y_%m_%d_%H_%M_%S')
        save_location = self.scenario.storage["location"]
        if not save_location or save_location in ["default", "auto", "topic_store"]:
            save_location = pathlib.Path(os.path.expanduser("~/.ros/topic_store/filesystem"))
        elif save_location.startswith("pkg="):
            package_name = save_location.split('=')[-1]
            save_location = pathlib.Path(os.path.expanduser("~/.ros/topic_store/filesystem/{}/".format(package_name)))
        else:
            save_location = pathlib.Path(os.path.expanduser(save_location))
        save_folder = save_location / self.scenario.context
        self.log("Configured save_location as '{}'".format(save_folder))
        try:
            save_folder.mkdir(parents=True)
        except OSError as e:
            if e.errno != 17:  # File exists is okay
                raise

        save_file = save_folder / "{}{}".format(formatted_datetime, TopicStorage.suffix)
        self.filesystem_storage = TopicStorage(save_file)
        self.log("Initialised saving to the filesystem at '{}'".format(self.filesystem_storage.path))

    def save_database(self, message_tree):
        insert_result = self.db_client.insert_one(message_tree)
        self.log("Inserted document to database result='acknowledged={}, inserted_id={}'".format(
            insert_result.acknowledged, insert_result.inserted_id))

    def save_filesystem(self, message_tree):
        self.log("Saving documents to file system n={}".format(self.saved_n))
        self.saved_n += 1
        self.filesystem_storage.insert_one(message_tree)

    def save(self):
        """Collates data from the scenario topic structure and saves. Returns SaveSuccess, SaveMessage"""
        data = self.subscriber_tree.get_message_tree()
        try:
            self.save_method_function(data)
        except Exception as e:
            self.log("Exception raised when saving! '{}'".format(e.message))
            return False, e.message
        return True, "Success!"
Ejemplo n.º 12
0
def _convert_cli():
    parser = argparse.ArgumentParser()
    parser.add_argument("-i",
                        "--input",
                        help="Input File",
                        type=str,
                        required=True)
    parser.add_argument("-o",
                        "--output",
                        help="Output File",
                        type=str,
                        required=True)
    parser.add_argument(
        "-c",
        "--collection",
        help="MongoDB collection to use if URI passed as --input",
        type=str,
        required=False)
    parser.add_argument("-q",
                        "--query",
                        help='MongoDB input query as dict (example: -q '
                        '\'{"_id": "ObjectId(5f718a354e5e8239dcd1eca1)"}\'',
                        type=str,
                        required=False,
                        default=None)
    parser.add_argument(
        "-p",
        "--projection",
        help='MongoDB input projection as dict (example: -p \'{"name": 1}\'',
        type=str,
        required=False,
        default=None)
    args = parser.parse_args()

    rospy.init_node("topic_store_convert", anonymous=True)
    input_path = pathlib.Path(args.input)
    output_path = pathlib.Path(args.output)

    # if not input_path.exists():
    #     raise IOError("Input file '{}' does not exist".format(input_path))

    if input_path.suffix == ".bag":
        raise NotImplementedError(
            "Converting from ROS bags is not currently supported. "
            "The conversion to ROS bags is lossy and requires adding meta data to reconstruct"
            "the original .topic_store or database documents")
    elif input_path.suffix == TopicStorage.suffix and output_path.suffix == ".bag":
        topic_store_to_ros_bag(input_path, output_path)
    elif input_path.suffix == ".yaml" and output_path.suffix == TopicStorage.suffix:
        input_path = resolve_scenario_yaml(input_path)
        mongodb_to_topic_store(MongoStorage.load(input_path), output_path)
    elif input_path.suffix == ".yaml" and output_path.suffix == ".bag":
        input_path = resolve_scenario_yaml(input_path)
        mongodb_to_ros_bag(MongoStorage.load(input_path), output_path)
    elif input_path.suffix == TopicStorage.suffix and output_path.suffix == ".yaml":
        output_path = resolve_scenario_yaml(output_path)
        topic_store_to_mongodb(input_path, output_path)
    elif is_uri(args.input):
        srv = args.input
        collection = args.collection
        query = args.query
        projection = args.projection

        # if not hasattr(args, "query") or not args.query:
        #     raise ValueError("If input is a MongoDB URI you must specify a DB query -q/--query to export data")
        if not hasattr(args, "collection") or not args.collection:
            raise ValueError(
                "If input is a MongoDB URI you must specify a DB collection -c/--collection to query data"
            )

        # Try to parse a query/projection string to a dict and perform some basic cleaning
        # The query string will filter the db documents by client.find(query)
        if query is not None:
            try:
                query, projection = [
                    x if x is None else json.loads(x)
                    for x in [args.query, args.projection]
                ]
            except ValueError:
                print(
                    "Query/Projection parameter cannot be parsed as a python dict \nQ: '{}'\nP: '{}'"
                    .format(args.query, args.projection))
                raise

        # Some simple rules to support searching by ID from console
        if query:
            for k, v in query.items():
                try:
                    unicode
                except NameError:
                    unicode = str
                if isinstance(v, (str, unicode)) and (v.startswith('ObjectId(')
                                                      and v.endswith(')')):
                    print("Converting query field '{}' to ObjectId".format(k))
                    query[k] = ObjectId(str(v[9:-1]))

        client = client_from_uri(srv, collection=collection)

        if is_uri(args.output):
            db2 = client_from_uri(args.output, collection=collection)
            mongodb_to_mongodb(client, db2, query=query, projection=projection)
        elif output_path.suffix == ".bag":
            mongodb_to_ros_bag(client,
                               output_path,
                               query=query,
                               projection=projection)
        elif output_path.suffix == TopicStorage.suffix:
            mongodb_to_topic_store(client,
                                   output_path,
                                   query=query,
                                   projection=projection)
        else:
            raise ValueError(
                "No valid conversion from Mongo URI '{}' to '{}' file".format(
                    client.uri, output_path))
    elif input_path.suffix == output_path.suffix:
        print("No conversion or migration for '{}' to '{}'".format(
            input_path, output_path))
        print(
            "If you would like to copy the file please use 'cp {} {}'".format(
                input_path, output_path))
    else:
        print("No conversion or migration for '{}' to '{}'".format(
            input_path, output_path))