Fix pathlib.Path support for FileSourceStage

morpheus/_lib/common/__init__.pyi

@@ -7,6 +7,7 @@
 from __future__ import annotations
 import morpheus._lib.common
 import typing
+import os
 
 __all__ = [
     "FiberQueue",
@@ -190,6 +191,10 @@ class TypeId():
     UINT8: morpheus._lib.common.TypeId # value = <TypeId.UINT8: 5>
     __members__: dict # value = {'EMPTY': <TypeId.EMPTY: 0>, 'INT8': <TypeId.INT8: 1>, 'INT16': <TypeId.INT16: 2>, 'INT32': <TypeId.INT32: 3>, 'INT64': <TypeId.INT64: 4>, 'UINT8': <TypeId.UINT8: 5>, 'UINT16': <TypeId.UINT16: 6>, 'UINT32': <TypeId.UINT32: 7>, 'UINT64': <TypeId.UINT64: 8>, 'FLOAT32': <TypeId.FLOAT32: 9>, 'FLOAT64': <TypeId.FLOAT64: 10>, 'BOOL8': <TypeId.BOOL8: 11>, 'STRING': <TypeId.STRING: 12>}
     pass
+@typing.overload
+def determine_file_type(filename: os.PathLike) -> FileTypes:
+    pass
+@typing.overload
 def determine_file_type(filename: str) -> FileTypes:
     pass
 def read_file_to_df(filename: str, file_type: FileTypes = FileTypes.Auto) -> object:

morpheus/_lib/common/module.cpp

@@ -37,7 +37,10 @@
 #include <pybind11/attr.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/pytypes.h>  // for return_value_policy::reference
+// for pathlib.Path -> std::filesystem::path conversions
+#include <pybind11/stl/filesystem.h>  // IWYU pragma: keep
 
+#include <filesystem>  // for std::filesystem::path
 #include <memory>
 #include <sstream>
 #include <string>
@@ -58,13 +61,29 @@ PYBIND11_MODULE(common, _module)
     CudfHelper::load();
 
     LoaderRegistry::register_factory_fn(
-        "file", [](nlohmann::json config) { return std::make_unique<FileDataLoader>(config); }, false);
+        "file",
+        [](nlohmann::json config) {
+            return std::make_unique<FileDataLoader>(config);
+        },
+        false);
     LoaderRegistry::register_factory_fn(
-        "grpc", [](nlohmann::json config) { return std::make_unique<GRPCDataLoader>(config); }, false);
+        "grpc",
+        [](nlohmann::json config) {
+            return std::make_unique<GRPCDataLoader>(config);
+        },
+        false);
     LoaderRegistry::register_factory_fn(
-        "payload", [](nlohmann::json config) { return std::make_unique<PayloadDataLoader>(config); }, false);
+        "payload",
+        [](nlohmann::json config) {
+            return std::make_unique<PayloadDataLoader>(config);
+        },
+        false);
     LoaderRegistry::register_factory_fn(
-        "rest", [](nlohmann::json config) { return std::make_unique<RESTDataLoader>(config); }, false);
+        "rest",
+        [](nlohmann::json config) {
+            return std::make_unique<RESTDataLoader>(config);
+        },
+        false);
 
     py::class_<TensorObject>(_module, "Tensor")
         .def_property_readonly("__cuda_array_interface__", &TensorObjectInterfaceProxy::cuda_array_interface)
@@ -106,9 +125,15 @@ PYBIND11_MODULE(common, _module)
         .value("CSV", FileTypes::CSV)
         .value("PARQUET", FileTypes::PARQUET);
 
-    _module.def("typeid_to_numpy_str", [](TypeId tid) { return DType(tid).type_str(); });
+    _module.def("typeid_to_numpy_str", [](TypeId tid) {
+        return DType(tid).type_str();
+    });
 
-    _module.def("determine_file_type", &determine_file_type, py::arg("filename"));
+    _module.def(
+        "determine_file_type", py::overload_cast<const std::string&>(&determine_file_type), py::arg("filename"));
+    _module.def("determine_file_type",
+                py::overload_cast<const std::filesystem::path&>(&determine_file_type),
+                py::arg("filename"));
     _module.def("read_file_to_df", &read_file_to_df, py::arg("filename"), py::arg("file_type") = FileTypes::Auto);
     _module.def("write_df_to_file",
                 &SerializersProxy::write_df_to_file,

morpheus/_lib/include/morpheus/objects/file_types.hpp

@@ -18,6 +18,7 @@
 #pragma once
 
 #include <cstdint>
+#include <filesystem>  // for path
 #include <ostream>
 #include <stdexcept>
 #include <string>
@@ -84,6 +85,15 @@ static inline std::ostream& operator<<(std::ostream& os, const FileTypes& f)
  */
 FileTypes determine_file_type(const std::string& filename);
 
+/**
+ * @brief Determines the file type from a filename based on extension. For example, my_file.json would return
+ * `FileTypes::JSON`.
+ *
+ * @param filename path to a file. Does not need to exist
+ * @return FileTypes
+ */
+FileTypes determine_file_type(const std::filesystem::path& filename);
+
 #pragma GCC visibility pop
 
 /** @} */  // end of group

morpheus/_lib/include/morpheus/stages/file_source.hpp

@@ -20,23 +20,17 @@
 #include "morpheus/messages/meta.hpp"
 
 #include <boost/fiber/context.hpp>
-#include <boost/fiber/future/future.hpp>
-#include <mrc/node/rx_sink_base.hpp>
-#include <mrc/node/rx_source_base.hpp>
-#include <mrc/node/source_properties.hpp>
 #include <mrc/segment/builder.hpp>
 #include <mrc/segment/object.hpp>
-#include <mrc/types.hpp>
 #include <pybind11/pytypes.h>
 #include <pymrc/node.hpp>
 #include <rxcpp/rx.hpp>  // for apply, make_subscriber, observable_member, is_on_error<>::not_void, is_on_next_of<>::not_void, trace_activity
 
-#include <map>
+#include <filesystem>  // for path
 #include <memory>
 #include <optional>
 #include <string>
 #include <thread>
-#include <vector>  // for vector
 
 namespace morpheus {
 /****** Component public implementations *******************/
@@ -98,6 +92,11 @@ struct FileSourceStageInterfaceProxy
                                                                        std::string filename,
                                                                        int repeat                   = 1,
                                                                        pybind11::dict parser_kwargs = pybind11::dict());
+    static std::shared_ptr<mrc::segment::Object<FileSourceStage>> init(mrc::segment::Builder& builder,
+                                                                       const std::string& name,
+                                                                       std::filesystem::path filename,
+                                                                       int repeat                   = 1,
+                                                                       pybind11::dict parser_kwargs = pybind11::dict());
 };
 #pragma GCC visibility pop
 /** @} */  // end of group

morpheus/_lib/src/objects/file_types.cpp

@@ -49,4 +49,9 @@ FileTypes determine_file_type(const std::string& filename)
     }
 }
 
+FileTypes determine_file_type(const std::filesystem::path& filename)
+{
+    return determine_file_type(filename.string());
+}
+
 }  // namespace morpheus

morpheus/_lib/src/stages/file_source.cpp

@@ -17,11 +17,7 @@
 
 #include "morpheus/stages/file_source.hpp"
 
-#include "mrc/node/rx_sink_base.hpp"
-#include "mrc/node/rx_source_base.hpp"
-#include "mrc/node/source_properties.hpp"
 #include "mrc/segment/object.hpp"
-#include "mrc/types.hpp"
 #include "pymrc/node.hpp"
 
 #include "morpheus/io/deserializers.hpp"
@@ -32,12 +28,12 @@
 #include <cudf/types.hpp>
 #include <glog/logging.h>
 #include <mrc/segment/builder.hpp>
-#include <pybind11/cast.h>
+#include <pybind11/cast.h>  // IWYU pragma: keep
 #include <pybind11/gil.h>
 #include <pybind11/pybind11.h>  // for str_attr_accessor
 #include <pybind11/pytypes.h>   // for pybind11::int_
 
-#include <functional>
+#include <filesystem>
 #include <memory>
 #include <optional>
 #include <sstream>
@@ -133,4 +129,14 @@ std::shared_ptr<mrc::segment::Object<FileSourceStage>> FileSourceStageInterfaceP
 
     return stage;
 }
+
+std::shared_ptr<mrc::segment::Object<FileSourceStage>> FileSourceStageInterfaceProxy::init(
+    mrc::segment::Builder& builder,
+    const std::string& name,
+    std::filesystem::path filename,
+    int repeat,
+    pybind11::dict parser_kwargs)
+{
+    return init(builder, name, filename.string(), repeat, std::move(parser_kwargs));
+}
 }  // namespace morpheus

morpheus/_lib/stages/__init__.pyi

@@ -11,6 +11,7 @@ import typing
 from morpheus._lib.common import FilterSource
 import morpheus._lib.common
 import mrc.core.segment
+import os
 
 __all__ = [
     "AddClassificationsStage",
@@ -45,6 +46,9 @@ class DeserializeMultiMessageStage(mrc.core.segment.SegmentObject):
     def __init__(self, builder: mrc.core.segment.Builder, name: str, batch_size: int, ensure_sliceable_index: bool = True) -> None: ...
     pass
 class FileSourceStage(mrc.core.segment.SegmentObject):
+    @typing.overload
+    def __init__(self, builder: mrc.core.segment.Builder, name: str, filename: os.PathLike, repeat: int, parser_kwargs: dict) -> None: ...
+    @typing.overload
     def __init__(self, builder: mrc.core.segment.Builder, name: str, filename: str, repeat: int, parser_kwargs: dict) -> None: ...
     pass
 class FilterDetectionsStage(mrc.core.segment.SegmentObject):

morpheus/_lib/stages/module.cpp

@@ -36,16 +36,21 @@
 #include "morpheus/utilities/http_server.hpp"  // for DefaultMaxPayloadSize
 #include "morpheus/version.hpp"
 
+#include <mrc/segment/builder.hpp>  // for Builder
 #include <mrc/segment/object.hpp>
 #include <mrc/utils/string_utils.hpp>
 #include <pybind11/attr.h>      // for multiple_inheritance
 #include <pybind11/pybind11.h>  // for arg, init, class_, module_, str_attr_accessor, PYBIND11_MODULE, pybind11
 #include <pybind11/pytypes.h>   // for dict, sequence
-#include <pymrc/utils.hpp>      // for pymrc::import
+// for pathlib.Path -> std::filesystem::path conversions
+#include <pybind11/stl/filesystem.h>  // IWYU pragma: keep
+#include <pymrc/utils.hpp>            // for pymrc::import
 #include <rxcpp/rx.hpp>
 
+#include <filesystem>  // for std::filesystem::path
 #include <memory>
 #include <sstream>
+#include <string>
 
 namespace morpheus {
 namespace py = pybind11;
@@ -107,11 +112,21 @@ PYBIND11_MODULE(stages, _module)
              py::arg("task_type")              = py::none(),
              py::arg("task_payload")           = py::none());
 
+    // py::overload_cast<const std::filesystem::path&>
     py::class_<mrc::segment::Object<FileSourceStage>,
                mrc::segment::ObjectProperties,
                std::shared_ptr<mrc::segment::Object<FileSourceStage>>>(
         _module, "FileSourceStage", py::multiple_inheritance())
-        .def(py::init<>(&FileSourceStageInterfaceProxy::init),
+        .def(py::init(py::overload_cast<mrc::segment::Builder&, const std::string&, std::string, int, py::dict>(
+                 &FileSourceStageInterfaceProxy::init)),
+             py::arg("builder"),
+             py::arg("name"),
+             py::arg("filename"),
+             py::arg("repeat"),
+             py::arg("parser_kwargs"))
+        .def(py::init(
+                 py::overload_cast<mrc::segment::Builder&, const std::string&, std::filesystem::path, int, py::dict>(
+                     &FileSourceStageInterfaceProxy::init)),
              py::arg("builder"),
              py::arg("name"),
              py::arg("filename"),

tests/common/test_determine_file_type.py

@@ -0,0 +1,32 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.
+#
+# 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 pathlib
+
+import pytest
+
+from morpheus.common import FileTypes
+from morpheus.common import determine_file_type
+
+
+@pytest.mark.parametrize("use_pathlib", [False, True])
+@pytest.mark.parametrize("ext, expected_result",
+                         [("csv", FileTypes.CSV), ("json", FileTypes.JSON), ("jsonlines", FileTypes.JSON),
+                          ("parquet", FileTypes.PARQUET)])
+def test_determine_file_type(ext: str, expected_result: FileTypes, use_pathlib: bool):
+    file_path = f"test.{ext}"
+    if use_pathlib:
+        file_path = pathlib.Path(file_path)
+
+    assert determine_file_type(file_path) == expected_result

tests/test_file_in_out.py

@@ -16,6 +16,8 @@
 
 import filecmp
 import os
+import pathlib
+import typing
 
 import numpy as np
 import pytest
@@ -24,6 +26,7 @@
 from _utils import assert_path_exists
 from _utils.dataset_manager import DatasetManager
 from morpheus.common import FileTypes
+from morpheus.config import Config
 from morpheus.config import CppConfig
 from morpheus.io.deserializers import read_file_to_df
 from morpheus.io.serializers import write_df_to_file
@@ -39,11 +42,22 @@
 
 @pytest.mark.slow
 @pytest.mark.parametrize("input_type", ["csv", "jsonlines", "parquet"])
+@pytest.mark.parametrize("use_pathlib", [False, True])
 @pytest.mark.parametrize("output_type", ["csv", "json", "jsonlines"])
 @pytest.mark.parametrize("flush", [False, True], ids=["no_flush", "flush"])
 @pytest.mark.parametrize("repeat", [1, 2, 5], ids=["repeat1", "repeat2", "repeat5"])
-def test_file_rw_pipe(tmp_path, config, input_type, output_type, flush, repeat: int):
+def test_file_rw_pipe(tmp_path: pathlib.Path,
+                      config: Config,
+                      input_type: str,
+                      use_pathlib: bool,
+                      output_type: str,
+                      flush: bool,
+                      repeat: int):
     input_file = os.path.join(TEST_DIRS.tests_data_dir, f'filter_probs.{input_type}')
+
+    if use_pathlib:
+        input_file = pathlib.Path(input_file)
+
     validation_file = os.path.join(TEST_DIRS.tests_data_dir, "filter_probs.csv")
     out_file = os.path.join(tmp_path, f'results.{output_type}')
 
@@ -77,7 +91,7 @@ def test_file_rw_pipe(tmp_path, config, input_type, output_type, flush, repeat:
     assert output_data.tolist() == validation_data.tolist()
 
 
-def test_file_read_json(config):
+def test_file_read_json(config: Config):
     src_file = os.path.join(TEST_DIRS.tests_data_dir, "simple.json")
 
     pipe = LinearPipeline(config)
@@ -98,7 +112,7 @@ def test_file_read_json(config):
 @pytest.mark.slow
 @pytest.mark.use_python
 @pytest.mark.usefixtures("chdir_tmpdir")
-def test_to_file_no_path(tmp_path, config):
+def test_to_file_no_path(tmp_path: pathlib.Path, config: Config):
     """
     Test to ensure issue #48 is fixed
     """
@@ -119,7 +133,7 @@ def test_to_file_no_path(tmp_path, config):
 @pytest.mark.slow
 @pytest.mark.parametrize("input_type", ["csv", "jsonlines", "parquet"])
 @pytest.mark.parametrize("output_type", ["csv", "json", "jsonlines"])
-def test_file_rw_multi_segment_pipe(tmp_path, config, input_type, output_type):
+def test_file_rw_multi_segment_pipe(tmp_path: pathlib.Path, config: Config, input_type: str, output_type: str):
     input_file = os.path.join(TEST_DIRS.tests_data_dir, f'filter_probs.{input_type}')
     validation_file = os.path.join(TEST_DIRS.tests_data_dir, "filter_probs.csv")
     out_file = os.path.join(tmp_path, f'results.{output_type}')
@@ -156,7 +170,7 @@ def test_file_rw_multi_segment_pipe(tmp_path, config, input_type, output_type):
                              os.path.join(TEST_DIRS.tests_data_dir, "filter_probs.csv"),
                              os.path.join(TEST_DIRS.tests_data_dir, "filter_probs_w_id_col.csv")
                          ])
-def test_file_rw_index_pipe(tmp_path, config, input_file):
+def test_file_rw_index_pipe(tmp_path: pathlib.Path, config: Config, input_file: str):
     out_file = os.path.join(tmp_path, 'results.csv')
 
     pipe = LinearPipeline(config)
@@ -183,7 +197,7 @@ def test_file_rw_index_pipe(tmp_path, config, input_file):
                          }), (os.path.join(TEST_DIRS.tests_data_dir, "filter_probs.jsonlines"), {})],
                          ids=["CSV", "CSV_ID", "JSON"])
 @pytest.mark.usefixtures("use_cpp")
-def test_file_roundtrip(tmp_path, input_file, extra_kwargs):
+def test_file_roundtrip(tmp_path: pathlib.Path, input_file: str, extra_kwargs: dict[str, typing.Any]):
 
     # Output file should be same type as input
     out_file = os.path.join(tmp_path, f'results{os.path.splitext(input_file)[1]}')
@@ -222,7 +236,7 @@ def test_read_cpp_compare(input_file: str):
 
 @pytest.mark.slow
 @pytest.mark.parametrize("output_type", ["csv", "json", "jsonlines"])
-def test_file_rw_serialize_deserialize_pipe(tmp_path, config, output_type):
+def test_file_rw_serialize_deserialize_pipe(tmp_path: pathlib.Path, config: Config, output_type: str):
     input_file = os.path.join(TEST_DIRS.tests_data_dir, "filter_probs.csv")
     out_file = os.path.join(tmp_path, f'results.{output_type}')
 
@@ -252,7 +266,7 @@ def test_file_rw_serialize_deserialize_pipe(tmp_path, config, output_type):
 
 @pytest.mark.slow
 @pytest.mark.parametrize("output_type", ["csv", "json", "jsonlines"])
-def test_file_rw_serialize_deserialize_multi_segment_pipe(tmp_path, config, output_type):
+def test_file_rw_serialize_deserialize_multi_segment_pipe(tmp_path: pathlib.Path, config: Config, output_type: str):
     input_file = os.path.join(TEST_DIRS.tests_data_dir, "filter_probs.csv")
     out_file = os.path.join(tmp_path, f'results.{output_type}')