From 8604e6282d88228796937c9e9692893c0a4fb3c6 Mon Sep 17 00:00:00 2001
From: Qi Zhao <zhaoqi99@outlook.com>
Date: Tue, 26 Mar 2024 16:01:36 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=F0=9F=90=9B=20loader=20loader=5Fextensi?=
 =?UTF-8?q?on=20path=20in=20Windows?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pyencrypt/encrypt.py  | 15 +++++++++++++--
 tests/test_encrypt.py |  6 +++++-
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/pyencrypt/encrypt.py b/pyencrypt/encrypt.py
index af3a0b7..1c328d3 100644
--- a/pyencrypt/encrypt.py
+++ b/pyencrypt/encrypt.py
@@ -1,4 +1,5 @@
 import os
+import sys
 import re
 from pathlib import Path
 from typing import Optional
@@ -97,8 +98,18 @@ def generate_so_file(cipher_key: str, d: int, n: int, base_dir: Optional[Path] =
         script_args=['build_ext', '--build-lib', temp_dir.as_posix()],
         cmdclass={'build_ext': build_ext},
     )
-    print(list(temp_dir.iterdir()))
-    return list(temp_dir.glob('loader.cpython-*-*.so'))[0].absolute()
+    if sys.platform.startswith('win'):
+        # loader.cp36-win_amd64.pyd
+        pattern = 'loader.cp*-*.pyd'
+    else:
+        # loader.cpython-36m-x86_64-linux-gnu.so
+        # loader.cpython-36m-darwin.so
+        pattern = "loader.cpython-*-*.so"
+
+    loader_extension = next(temp_dir.glob(pattern), None)
+    if loader_extension is None:
+        raise Exception(f"Can't find loader extension in {temp_dir.as_posix()}")
+    return loader_extension.absolute()
 
 
 def encrypt_file(path: Path, key: str, delete_origin: bool = False, new_path: Optional[Path] = None):
diff --git a/tests/test_encrypt.py b/tests/test_encrypt.py
index f5f7736..c0e51cc 100644
--- a/tests/test_encrypt.py
+++ b/tests/test_encrypt.py
@@ -1,6 +1,7 @@
 import os
 from pathlib import Path
 import shutil
+import sys
 
 import pytest
 from pyencrypt.encrypt import can_encrypt, encrypt_file, encrypt_key, generate_so_file
@@ -58,7 +59,10 @@ def test_generate_so_file_default_path(self, key):
         assert generate_so_file(cipher_key, d, n)
         assert (Path(os.getcwd()) / 'encrypted' / 'loader.py').exists() is True
         assert (Path(os.getcwd()) / 'encrypted' / 'loader_origin.py').exists() is True
-        assert list((Path(os.getcwd()) / 'encrypted').glob('loader.cpython-*-*.so')) != []
+        if sys.platform.startswith('win'):
+            assert next((Path(os.getcwd()) / 'encrypted').glob('loader.cp*-*.pyd'), None) is not None
+        else:
+            assert next((Path(os.getcwd()) / 'encrypted').glob('loader.cpython-*-*.so'), None) is not None
 
 
 @pytest.mark.parametrize(