From dd88065a4ba07110e7f3ae05f09c6847e78b0d54 Mon Sep 17 00:00:00 2001
From: Tommy van der Vorst <tommy@pixelspark.nl>
Date: Sun, 5 Mar 2023 20:17:15 +0100
Subject: [PATCH] fix(!): rename PySession -> Session

---
 README.md                           | 4 ++--
 wonnx-py/src/lib.rs                 | 2 +-
 wonnx-py/tests/test_onnx_backend.py | 2 +-
 wonnx-py/tests/test_wonnx.py        | 8 ++++----
 4 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index 568cbd54..4aa011e6 100644
--- a/README.md
+++ b/README.md
@@ -48,8 +48,8 @@ pip install wonnx
 And then, to use:
 
 ```python
-from wonnx import PySession
-session = PySession.from_path(
+from wonnx import Session
+session = Session.from_path(
     "../data/models/single_relu.onnx"
 )
 inputs = {"x": [-1.0, 2.0]}
diff --git a/wonnx-py/src/lib.rs b/wonnx-py/src/lib.rs
index 4fb99e8f..eb9af8dd 100644
--- a/wonnx-py/src/lib.rs
+++ b/wonnx-py/src/lib.rs
@@ -4,7 +4,7 @@ use pyo3::types::PyDict;
 use std::collections::HashMap;
 
 use ::wonnx::Session;
-#[pyclass]
+#[pyclass(name = "Session")]
 #[repr(transparent)]
 pub struct PySession {
     pub session: Session,
diff --git a/wonnx-py/tests/test_onnx_backend.py b/wonnx-py/tests/test_onnx_backend.py
index 96d9d154..af51913f 100644
--- a/wonnx-py/tests/test_onnx_backend.py
+++ b/wonnx-py/tests/test_onnx_backend.py
@@ -37,7 +37,7 @@ def __init__(self, inputs, outputs, outputs_shape, model):
         self.inputs = inputs
         self.outputs = outputs
         self.outputs_shape = outputs_shape
-        self.session = wonnx.PySession.from_bytes(onnx._serialize(model))
+        self.session = wonnx.Session.from_bytes(onnx._serialize(model))
         self.rtol = 1
         pass
 
diff --git a/wonnx-py/tests/test_wonnx.py b/wonnx-py/tests/test_wonnx.py
index 6e60badc..5191704e 100644
--- a/wonnx-py/tests/test_wonnx.py
+++ b/wonnx-py/tests/test_wonnx.py
@@ -28,7 +28,7 @@ def test_parse_model():
     model_def = helper.make_model(graph_def, producer_name="onnx-example")
     model = onnx.shape_inference.infer_shapes(model_def)
 
-    session = wonnx.PySession.from_bytes(onnx._serialize(model))
+    session = wonnx.Session.from_bytes(onnx._serialize(model))
     inputs = {"x": [-1.0, 2.0]}
     assert session.run(inputs) == {"y": [0.0, 2.0]}, "Single Relu does not work"
 
@@ -36,7 +36,7 @@ def test_parse_model():
 def test_from_path():
 
     # Create the model (ModelProto)
-    session = wonnx.PySession.from_path(
+    session = wonnx.Session.from_path(
         "../examples/data/models/single_relu.onnx"
     )
     inputs = {"x": [-1.0, 2.0]}
@@ -51,7 +51,7 @@ def test_mnist():
     input = np.reshape(gray, (1, 1, 28, 28))
     # Create the model (ModelProto)
 
-    session = wonnx.PySession.from_path(
+    session = wonnx.Session.from_path(
         "../examples/data/models/opt-mnist.onnx"
     )
     inputs = {"Input3": input.flatten().tolist()}
@@ -79,7 +79,7 @@ def test_squeezenet():
     input_tensor = transform(image)
     # Create the model (ModelProto)
 
-    session = wonnx.PySession.from_path(
+    session = wonnx.Session.from_path(
         "../examples/data/models/opt-squeeze.onnx"
     )
     inputs = {"data": input_tensor.flatten().tolist()}