v0.1.2

Author: swoiow

.github/scripts/repack_wheels.py

@@ -0,0 +1,74 @@
+import shutil
+import subprocess
+import sys
+import tempfile
+import zipfile
+from pathlib import Path
+
+
+def _read_version(cargo_toml: Path) -> str:
+    for line in cargo_toml.read_text(encoding="utf-8").splitlines():
+        if line.strip().startswith("version"):
+            return line.split("=", 1)[1].strip().strip('"')
+    raise SystemExit("Version not found in Cargo.toml")
+
+
+def _load_setup_template(template_path: Path, version: str) -> str:
+    template = template_path.read_text(encoding="utf-8")
+    if "__VERSION__" not in template:
+        raise SystemExit("setup.py template missing __VERSION__ placeholder")
+    return template.replace("__VERSION__", version)
+
+
+def main() -> None:
+    repo_root = Path.cwd()
+    dist_dir = repo_root / "dist"
+    cargo_toml = repo_root / "Cargo.toml"
+    build_src = repo_root / "build_bin.py"
+    ipsearcher_src = repo_root / "example_production.py"
+    setup_template = repo_root / ".github" / "scripts" / "setup.py"
+
+    if not cargo_toml.exists():
+        raise SystemExit("Cargo.toml not found")
+    if not build_src.exists():
+        raise SystemExit("build_bin.py not found")
+    if not ipsearcher_src.exists():
+        raise SystemExit("example_production.py not found")
+    if not setup_template.exists():
+        raise SystemExit("setup.py template not found")
+
+    version = _read_version(cargo_toml)
+    setup_py = _load_setup_template(setup_template, version)
+
+    wheels = sorted(dist_dir.glob("*.whl"))
+    if not wheels:
+        raise SystemExit("No wheels found in dist")
+
+    for wheel in wheels:
+        temp_dir = Path(tempfile.mkdtemp())
+        with zipfile.ZipFile(wheel, "r") as zf:
+            zf.extractall(temp_dir)
+        package_dir = temp_dir / "poptrie"
+        if not package_dir.exists():
+            raise SystemExit(f"poptrie package not found in {wheel.name}")
+
+        shutil.copy2(build_src, package_dir / "build.py")
+        shutil.copy2(ipsearcher_src, package_dir / "ipsearcher.py")
+        (temp_dir / "setup.py").write_text(setup_py, encoding="utf-8")
+
+        wheel.unlink()
+        subprocess.run(
+            [sys.executable, "setup.py", "bdist_wheel"],
+            cwd=temp_dir,
+            check=True,
+        )
+        built_wheels = sorted((temp_dir / "dist").glob("*.whl"))
+        if not built_wheels:
+            raise SystemExit("bdist_wheel did not produce a wheel")
+        for built_wheel in built_wheels:
+            shutil.copy2(built_wheel, dist_dir / built_wheel.name)
+        shutil.rmtree(temp_dir)
+
+
+if __name__ == "__main__":
+    main()

.github/scripts/setup.py

@@ -0,0 +1,12 @@
+from setuptools import find_packages, setup
+
+
+setup(
+    name="poptrie",
+    version="__VERSION__",
+    description="Fast IP lookup using poptrie",
+    author="HarmonSir",
+    author_email="[email protected]",
+    license="Apache-2.0",
+    packages=find_packages(),
+)

.github/workflows/release.yml

@@ -81,11 +81,15 @@ jobs:
           echo "PYTHON_EXEC=$(pwd)/.venv/Scripts/python.exe" >> $GITHUB_ENV
 
       - name: Install maturin
-        run: $PYTHON_EXEC -m pip install --upgrade pip maturin
+        run: $PYTHON_EXEC -m pip install --upgrade pip maturin setuptools wheel
 
       - name: Build wheels
         run: $PYTHON_EXEC -m maturin build --release --out dist
 
+      - name: Repack wheels
+        run: |
+          $PYTHON_EXEC .github/scripts/repack_wheels.py
+
       - name: Install wheel
         run: $PYTHON_EXEC -m pip install dist/*.whl
 
@@ -129,6 +133,31 @@ jobs:
                 throw error;
               }
             }
+      - name: Remove latest release
+        uses: actions/github-script@v7
+        with:
+          script: |
+            try {
+              const release = await github.rest.repos.getReleaseByTag({
+                owner: context.repo.owner,
+                repo: context.repo.repo,
+                tag: "latest",
+              });
+              await github.rest.repos.deleteRelease({
+                owner: context.repo.owner,
+                repo: context.repo.repo,
+                release_id: release.data.id,
+              });
+              await github.rest.git.deleteRef({
+                owner: context.repo.owner,
+                repo: context.repo.repo,
+                ref: "tags/latest",
+              });
+            } catch (error) {
+              if (error.status !== 404) {
+                throw error;
+              }
+            }
 
       - name: Download wheels
         uses: actions/download-artifact@v4

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "poptrie"
-version = "0.1.1"
+version = "0.1.2"
 edition = "2021"
 
 [lib]

build_bin.py

@@ -125,11 +125,12 @@ def save(self, output_path):
         print(f"构建完成！压缩后大小: {len(final_data) / 1024:.2f} KB")
 
 
-# --- 使用方式 ---
-builder = BinBuilder()
-# 这里放入你收集的中国 IP CIDR 列表
-china_cidrs = ["1.0.1.0/24", "110.16.0.0/12", "240e::/18"]
-for c in china_cidrs:
-    builder.add_cidr(c)
-
-builder.save("china_ip.bin")
+if __name__ == '__main__':
+    # --- 使用方式 ---
+    builder = BinBuilder()
+    # 这里放入你收集的中国 IP CIDR 列表
+    china_cidrs = ["1.0.1.0/24", "110.16.0.0/12", "240e::/18"]
+    for c in china_cidrs:
+        builder.add_cidr(c)
+
+    builder.save("china_ip.bin")

example_production.py

@@ -24,14 +24,14 @@ def __init__(self, bin_path: str | Path) -> None:
         """
         self._searcher = poptrie.IpSearcher(str(bin_path))
 
-    def is_china_ip(self, ip_bytes: bytes) -> bool:
+    def is_matched_ip(self, ip_bytes: bytes) -> bool:
         """Check a single IPv4/IPv6 address in packed bytes.
         查询单个 IP（IPv4/IPv6）对应的 bytes。
 
         :param ip_bytes: 4-byte IPv4 or 16-byte IPv6.
         :return: True if matched, otherwise False.
         """
-        return self._searcher.is_china_ip(ip_bytes)
+        return self._searcher.is_matched_ip(ip_bytes)
 
     def batch_check_strings(self, ips: Iterable[str]) -> list[bool]:
         """Check a list of IP strings, preserving input order.
@@ -131,7 +131,7 @@ def main() -> None:
     searcher = IpSearcher(bin_path)
 
     ip_bytes = socket.inet_pton(socket.AF_INET, "1.0.1.1")
-    print(searcher.is_china_ip(ip_bytes))
+    print(searcher.is_matched_ip(ip_bytes))
 
     ips = ["1.0.1.1", "8.8.8.8", "240e::1", "2001:db8::"]
     print(searcher.check_ips(ips))

src/lib.rs

@@ -24,7 +24,7 @@ impl IpSearcher {
     }
 
     /// 核心查询逻辑：支持 IPv4 (4字节) 和 IPv6 (16字节)
-    fn is_china_ip(&self, ip_bytes: &[u8]) -> bool {
+    fn is_matched_ip(&self, ip_bytes: &[u8]) -> bool {
         let mut cursor: usize = 0;
 
         for &byte in ip_bytes {
@@ -63,16 +63,16 @@ impl IpSearcher {
         // 这是极致性能的关键：内存完全连续，没有 Python 对象开销
         packed_ips
             .chunks_exact(ip_stride)
-            .map(|ip_chunk| self.is_china_ip(ip_chunk))
+            .map(|ip_chunk| self.is_matched_ip(ip_chunk))
             .collect()
     }
 
     fn batch_check_strings(&self, py: Python<'_>, ips: Vec<String>) -> Vec<bool> {
         py.allow_threads(|| {
             ips.into_par_iter()
                 .map(|ip_str| match ip_str.parse::<IpAddr>() {
-                    Ok(IpAddr::V4(v4)) => self.is_china_ip(&v4.octets()),
-                    Ok(IpAddr::V6(v6)) => self.is_china_ip(&v6.octets()),
+                    Ok(IpAddr::V4(v4)) => self.is_matched_ip(&v4.octets()),
+                    Ok(IpAddr::V6(v6)) => self.is_matched_ip(&v6.octets()),
                     Err(_) => false,
                 })
                 .collect()