Python wrappers (#110)

* wrap C++ objects in python

Author: bcebere

private_set_intersection/python/BUILD

@@ -19,20 +19,14 @@ pybind_extension(
 
 py_library(
     name = "openmined_psi",
-    srcs = ["__init__.py"],
+    srcs = [
+        "__init__.py",
+    ],
     data = ["//private_set_intersection/python:_openmined_psi.so"],
     srcs_version = "PY3",
     visibility = ["//visibility:public"],
 )
 
-py_binary(
-    name = "openmined_psi_bin",
-    srcs = ["__init__.py"],
-    data = ["//private_set_intersection/python:_openmined_psi.so"],
-    main = "__init__.py",
-    srcs_version = "PY3",
-)
-
 py_test(
     name = "tests",
     srcs = ["tests.py"],

private_set_intersection/python/__init__.py

@@ -1,21 +1,156 @@
 """Private Set Intersection protocol based on ECDH and Bloom
 Filters.
 """
+from typing import List
+
 try:
     # Used in Bazel envs
     from private_set_intersection.python import _openmined_psi as psi
 except ImportError:
     # Default package
-    import openmined_psi as psi
+    import _openmined_psi as psi
 
-client = psi.client
-server = psi.server
+__version__ = psi.__version__
 
-proto_server_setup = psi.proto_server_setup
-proto_request = psi.proto_request
-proto_response = psi.proto_response
 
-__version__ = psi.__version__
+proto_server_setup = psi.cpp_proto_server_setup
+proto_request = psi.cpp_proto_request
+proto_response = psi.cpp_proto_response
+
+
+class client:
+    def __init__(self, data: psi.cpp_client):
+        """Constructor method for the client object.
+        Args:
+            data: cpp_client object.
+        Returns:
+            client object.
+        """
+        self.data = data
+
+    @classmethod
+    def CreateWithNewKey(cls, reveal_intersection: bool):
+        """Constructor method for the client object.
+        Args:
+            reveal_intersection: indicates whether the client wants to learn the elements in the intersection or only its size.
+        Returns:
+            client object.
+        """
+        return cls(psi.cpp_client.CreateWithNewKey(reveal_intersection))
+
+    @classmethod
+    def CreateFromKey(cls, key_bytes: bytes, reveal_intersection: bool):
+        """Constructor method for the client object.
+        Args:
+            reveal_intersection: indicates whether the client wants to learn the elements in the intersection or only its size.
+            key_bytes: existing encryption key.
+        Returns:
+            client object.
+        """
+        return cls(psi.cpp_client.CreateFromKey(key_bytes, reveal_intersection))
+
+    def CreateRequest(self, data: List[str]) -> proto_request:
+        """Create a request protobuf to be serialized and sent to the server.
+        Args:
+            data: client items.
+        Returns:
+            A Protobuffer with the request.
+        """
+        return self.data.CreateRequest(data)
+
+    def GetIntersection(
+        self, server_setup: proto_server_setup, server_response: proto_response
+    ) -> List[int]:
+        """Process the server's response and return the intersection of the client and server inputs.
+        Args:
+            server_setup: A protobuffer with the setup message.
+            server_response: A protobuffer with server's response.
+        Returns:
+            A list of indices in clients set.
+        """
+        return self.data.GetIntersection(server_setup, server_response)
+
+    def GetIntersectionSize(
+        self, server_setup: proto_server_setup, server_response: proto_response
+    ) -> int:
+        """Process the server's response and return the size of the intersection.
+        Args:
+            server_setup: A protobuffer with the setup message.
+            server_response: A protobuffer with server's response.
+        Returns:
+            The intersection size.
+        """
+        return self.data.GetIntersectionSize(server_setup, server_response)
+
+    def GetPrivateKeyBytes(self) -> bytes:
+        """Returns this instance's private key. This key should only be used to create other client instances. DO NOT SEND THIS KEY TO ANY OTHER PARTY!
+        Returns:
+            Bytes containing the key.
+        """
+        return self.data.GetPrivateKeyBytes()
+
+
+class server:
+    def __init__(self, data: psi.cpp_server):
+        """Constructor method for the server object.
+        Args:
+            data: cpp_server object.
+        Returns:
+            server object.
+        """
+        self.data = data
+
+    @classmethod
+    def CreateWithNewKey(cls, reveal_intersection: bool):
+        """Constructor method for the server object.
+        Args:
+            reveal_intersection: indicates whether the server supports to return the elements in the intersection or only its size.
+        Returns:
+            server object.
+        """
+        return cls(psi.cpp_server.CreateWithNewKey(reveal_intersection))
+
+    @classmethod
+    def CreateFromKey(cls, key_bytes: bytes, reveal_intersection: bool):
+        """Constructor method for the server object.
+        Args:
+            reveal_intersection: indicates whether the server supports to return the elements in the intersection or only its size.
+            key_bytes: existing encryption key.
+        Returns:
+        Returns:
+            server object.
+        """
+        return cls(psi.cpp_server.CreateFromKey(key_bytes, reveal_intersection))
+
+    def CreateSetupMessage(
+        self, fpr: float, num_client_inputs: int, inputs: List[str]
+    ) -> proto_server_setup:
+        """Create a setup message from the server's dataset to be sent to the client.
+        Args:
+            fpr: the probability that any query of size `num_client_inputs` will result in a false positive.
+            num_client_inputs: Client set size.
+            inputs: Server items.
+        Returns:
+            A Protobuf with the setup message.
+        """
+        return self.data.CreateSetupMessage(fpr, num_client_inputs, inputs)
+
+    def ProcessRequest(self, client_request: proto_request) -> proto_response:
+        """Process a client query and returns the corresponding server response to be sent to the client.
+        Args:
+            client_request: A Protobuf containing the client request
+        Returns:
+            A Protobuf with the server response.
+        """
+        return self.data.ProcessRequest(client_request)
+
+    def GetPrivateKeyBytes(self) -> bytes:
+        """Return this instance's private key. This key should only be used to create other server instances. DO NOT SEND THIS KEY TO ANY OTHER PARTY!
+        Returns:
+            Bytes containing the key.
+        """
+        return self.data.GetPrivateKeyBytes()
+
 
 __all__ = [
     "client",

private_set_intersection/python/psi_bindings.cpp

@@ -40,7 +40,7 @@ void bind(pybind11::module& m) {
       "Filters";
 
   m.attr("__version__") = ::private_set_intersection::Package::kVersion;
-  py::class_<psi_proto::ServerSetup>(m, "proto_server_setup")
+  py::class_<psi_proto::ServerSetup>(m, "cpp_proto_server_setup")
       .def(py::init<>())
       .def(
           "bits",
@@ -64,7 +64,7 @@ void bind(pybind11::module& m) {
         loadProto(obj, data);
         return obj;
       });
-  py::class_<psi_proto::Request>(m, "proto_request")
+  py::class_<psi_proto::Request>(m, "cpp_proto_request")
       .def(py::init<>())
       .def("encrypted_elements_size",
            &psi_proto::Request::encrypted_elements_size)
@@ -99,7 +99,7 @@ void bind(pybind11::module& m) {
         loadProto(obj, data);
         return obj;
       });
-  py::class_<psi_proto::Response>(m, "proto_response")
+  py::class_<psi_proto::Response>(m, "cpp_proto_response")
       .def(py::init<>())
       .def("encrypted_elements_size",
            &psi_proto::Response::encrypted_elements_size)
@@ -131,7 +131,7 @@ void bind(pybind11::module& m) {
         return obj;
       });
 
-  py::class_<psi::PsiClient>(m, "client")
+  py::class_<psi::PsiClient>(m, "cpp_client")
       .def_static(
           "CreateWithNewKey",
           [](bool reveal_intersection) {
@@ -183,7 +183,7 @@ void bind(pybind11::module& m) {
           },
           py::call_guard<py::gil_scoped_release>());
 
-  py::class_<psi::PsiServer>(m, "server")
+  py::class_<psi::PsiServer>(m, "cpp_server")
       .def_static(
           "CreateWithNewKey",
           [](bool reveal_intersection) {

setup.py

@@ -84,16 +84,25 @@ def bazel_build(self, ext):
         ]
         self.spawn(bazel_argv)
 
+        ext.name = "_" + ext.name
         shared_lib_ext = ".so"
-        shared_lib = "_" + ext.name + shared_lib_ext
+        shared_lib = ext.name + shared_lib_ext
         ext_bazel_bin_path = os.path.join(self.build_temp, "bazel-bin", ext.relpath, shared_lib)
 
         ext_dest_path = self.get_ext_fullpath(ext.name)
         ext_dest_dir = os.path.dirname(ext_dest_path)
+
         if not os.path.exists(ext_dest_dir):
             os.makedirs(ext_dest_dir)
         shutil.copyfile(ext_bazel_bin_path, ext_dest_path)
 
+        package_dir = os.path.join(ext_dest_dir, "openmined_psi")
+        if not os.path.exists(package_dir):
+            os.makedirs(package_dir)
+        shutil.copyfile(
+            "private_set_intersection/python/__init__.py", os.path.join(package_dir, "__init__.py")
+        )
+
 
 setuptools.setup(
     name="openmined.psi",
@@ -103,7 +112,6 @@ def bazel_build(self, ext):
     url="https://github.com/OpenMined/PSI",
     python_requires=">=3.6",
     package_dir={"": "private_set_intersection/python"},
-    packages=setuptools.find_packages("private_set_intersection/python"),
     cmdclass=dict(build_ext=BuildBazelExtension),
     ext_modules=[
         BazelExtension("openmined_psi", "//private_set_intersection/python:openmined_psi",)