您的位置:  首页 > 技术 > python语言 > 正文

实战经验分享:使用 PyO3 来构建你的 Python 模块

2021-10-14 15:00 OSCHINA 又拍云 次阅读 条评论

PyO3 主要用于创建原生 Python 的扩展模块。PyO3 还支持从 Rust 二进制文件运行 Python 代码并与之交互,可以实现 rust 与 Python 代码共存。在一些对性能要求较高的模块上,可以考虑使用 PyO3 构建对应的功能模块。PyO3 的功能分离,不用过多担心模块之间的耦合性,并且在速度上能有一定的提升。

github地址: https://github.com/PyO3/pyo3

版本规定如下:

  • Python 3.6+

  • Rust 1.41+

接下来我们通过一个小的 demo 了解一下从 PyO3 编译模块到 Python 中正常使用的整个流程。

cargo new --lib string-sum

创建项目

# lib.rs

[package]

name = "string-sum"

version = "0.1.0"

edition = "2018"



[lib]

name = "string_sum"

# "cdylib" is necessary to produce a shared library for Python to import from.

#

# Downstream Rust code (including code in `bin/`, `examples/`, and `tests/`) will not be able

# to `use string_sum;` unless the "rlib" or "lib" crate type is also included, e.g.:

# crate-type = ["cdylib", "rlib"]

crate-type = ["cdylib"]



[dependencies.pyo3]

version = "0.14.1"

features = ["extension-module"] // 扩展模块,像其他的还有auto-initialize

// src/lib.rs

use std::usize;



use  pyo3::prelude::*;



// like this

// def sum_as_string(a:str, b:str) -> str:

//      return a+b

#[pyfunction]

fn sum_as_string(a: usize, b: usize) -> PyResult<String>{

    Ok((a+b).to_string())

}



// Mount method to module 

#[pymodule]

fn string_sum(py: Python, m: &PyModule) -> PyResult<()>{

    m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;

    Ok(())

}

编译与使用

编译完成之后,我们会在 target 文件夹下面发现一个 wheel 文件。文件名组合为 “模块名 + 当前 Python 版本+当前系统型号”,比如:string_sum-0.1.0-cp39-cp39-macosx_10_7_x86_64.whl

pip3 install ./target/wheel/string_sum-0.1.0-cp39-cp39-macosx_10_7_x86_64.whl

创建 python 文件:

# example.py

from string_sum import sum_as_string

print(sum_as_string(1,2))

# echo 3

编译工具的选择和使用

官方提供了两种编译工具的选择:

  • rust 写的 maturin

  • 传统的setup.py的方式

使用 maturin 编译

# 安装 

pip3 install maturin

# 编译

maturin build

# maturin publish 发布

# 虚拟环境中使用 会自动去寻找/target/wheel/ 下的 *.wheel文件然后安装

virtualenv venv

source ./venv/bin/activate

maturin develop

使用 setup.py 编译

# 安装

pip3 install setuptools-rust

编写 setup.py 文件:

# setup.py





from setuptools import setup

from setuptools_rust import Binding, RustExtension



setup(

    # 包名称

    name="string_sum", 

    # 包版本 

    version="0.1",

    # rust扩展 其中"string_sum.string_sum"中

    # 第一个string_sum 指的是当前的包

    # 第二个指的是

    # #[pymodule]

    # fn string_sum(py: Python, m: &PyModule) -> PyResult<()>{

    #     m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;

    #     Ok(())

    # }

    # 中的string_sum

    rust_extensions=[

        RustExtension(

            "string_sum.string_sum", 

            binding=Binding.PyO3,

            debug=False

            )

    ],

    # 需要创建一个文件夹 string_sum

    packages=["string_sum"],

    # rust extensions are not zip safe, just like C-extensions.

    zip_safe=False,

    # 标注

    classifiers=[

        "License :: OSI Approved :: MIT License",

        "Development Status :: 3 - Alpha",

        "Intended Audience :: Developers",

        "Programming Language :: Python",

        "Programming Language :: Rust",

        "Operating System :: POSIX",

        "Operating System :: MacOS :: MacOS X",

    ],

    include_package_data=True

)

# 打包

mkdir string_sum

touch string_sum/__init__.py

virtualenv venv && source venv/bin/activate

pip setup.py build && pip setup.py install && pip setup.py develop

会引用本地的文件:

docker 中的应用

同样的,如果创建的 App 本身是在 docker 内部运行的。那么第一步我们需要安装 rust 的环境 dockerfile。具体如下:

#!/bin/bash

curl https://sh.rustup.rs -sSf | bash -s -- -y

source $HOME/.cargo/env

rustc --version

python setup.py install

# ddockerfile 

FROM python:3.7

WORKDIR /app

ADD . /app

RUN pip install --upgrade pip \

    && pip install -r requirements.txt

RUN ./init.sh

CMD [python, xx.py]

# requirements.txt

semantic-version==2.8.5

setuptools-rust==0.12.1

toml==0.10.2

# rust国内镜像源 config

# /root/.cargo/config

[source.crates-io]

registry = "https://github.com/rust-lang/crates.io-index"

replace-with = 'ustc'

[source.ustc]

registry = "git://mirrors.ustc.edu.cn/crates.io-index"

[term]

verbose = true

color = 'auto'

具体目录如下:

-rw-r--r-- Cargo.lock

-rw-r--r-- Cargo.toml

-rw-r--r-- config           # 配置文件

-rw-r--r-- Dockerfile

-rwxrwxrwx init.sh          # 初始化rust环境脚本

-rw-r--r-- requirements.txt

-rw-r--r-- setup.py         # 打包脚本

drwxr-xr-x src              # rust项目

drwxr-xr-x string_sum 

-rw-r--r-- xx.py            # 可行性测试文件

用 PyO3 写一个 Python 的rsa加解密包

看过之前的文章的小伙伴《灵魂画手:漫画图解 SSH》 ,应该对 rsa 的整个加解密流程有所了解啦。那我们不妨用 PyO3 来构建一个 Python 的 rsa 加解密包。使用场景如下:

客户端本地生成公私钥,通过前期认证过程,将公钥发送给服务端保存,后期通信过程中,客户端主动发送消息给服务端,客户端通过私钥对信息加密,服务端通过对应的公钥进行解密。

github 地址: https://github.com/hzjsea/pyo3-crypto

后续又扩展了一些内容,比如 MD5 加密,签名等等。

# 自动化脚本

#!/bin/bash



echo "init......"



# set python version 

# INSTALL_PYTHON_VERSION=python3.6



find_python() {

        set +e

        unset BEST_VERSION

        for V in 37 3.7 38 3.8 39 3.9 3; do

                if which python$V >/dev/null; then

                        if [ "$BEST_VERSION" = "" ]; then

                                BEST_VERSION=$V

                        fi

                fi

        done

        echo $BEST_VERSION

        set -e

}



if [ "$INSTALL_PYTHON_VERSION" = "" ]; then

        INSTALL_PYTHON_VERSION=$(find_python)

fi



# This fancy syntax sets INSTALL_PYTHON_PATH to "python3.7", unless

# INSTALL_PYTHON_VERSION is defined.

# If INSTALL_PYTHON_VERSION equals 3.8, then INSTALL_PYTHON_PATH becomes python3.8

# 找不到就python3.7

INSTALL_PYTHON_PATH=python${INSTALL_PYTHON_VERSION:-3.7}

echo $INSTALL_PYTHON_PATH



echo "Python version is $INSTALL_PYTHON_VERSION"

$INSTALL_PYTHON_PATH -m venv venv

if [ ! -f "activate" ]; then

        ln -s venv/bin/activate .

fi



. ./activate



python -m pip install --upgrade pip

python -m pip install wheel

python -m pip install -r ./requirements.txt

maturin build

maturin develop



current_shell=$(echo $SHELL)

if current_shell=/bin/bash; then

    echo  "PASS: source /venv/bin/activate >> ~/.bashrc"

elif current_shell=/bin/zsh;then

    echo "PASS: source /venv/bin/activate >> ~/.zshrc"

fi

//  src/lib.rs 文件

use std::u32;

use pyo3::prelude::*;

use openssl::rsa::{Padding,Rsa};



const SECRET: &'static str = "CHFfxQA3tqEZgKusgwZjmI5lFsoZxXGXnQLA97oYga2M33sLwREZyy1mWCM8GIIA";



mod crypto_utils {

    use hmac::{Hmac, Mac, NewMac};

    use sha2::Sha256;

    use std::fmt::Write;



    type Hmacsha256 = Hmac<Sha256>;

    fn encode_hex(bytes: &[u8]) -> String {

        let mut s = String::with_capacity(bytes.len() * 2);

        for &b in bytes {

            match write!(&mut s, "{:02x}", b) {

                Ok(_) => {},

                Err(_) => {}

            };

        }

        s

    }



    pub fn hash_hmac(secret: &str, msg: &str) -> String {

        let mut mac = Hmacsha256::new_from_slice(secret.as_bytes()).expect("HMAC can take key of any size");

        mac.update(msg.as_bytes());

        let result = mac.finalize();

        let code_bytes = result.into_bytes();

        encode_hex(&code_bytes)

    }



}



// create public/private key  create_key(1024)

fn create_key(len:u32) -> (String,String){

    let rsa = openssl::rsa::Rsa::generate(len).unwrap();

    let pubkey = String::from_utf8(rsa.public_key_to_pem().unwrap()).unwrap();

    let prikey  = String::from_utf8(rsa.private_key_to_pem().unwrap()).unwrap();



    (pubkey, prikey)

}







#[pyclass]

struct Crypto {

    // #[pyo3(get, set)]

    // pubkey: String,

    // #[pyo3(get,set)]

    // prikey: String,

    pub_key: Rsa<openssl::pkey::Public>,

    pri_key: Rsa<openssl::pkey::Private>

}



#[pyfunction]

fn generate_key(len:u32) -> (String, String){

    create_key(len)

}



#[pymethods]

impl Crypto {

    #[new]

    pub fn __new__(pubkey: &str,prikey: &str) -> Self {

        Crypto {

            // pubkey: pubkey.to_owned(),

            // prikey: prikey.to_owned(),

            pub_key: Rsa::public_key_from_pem(pubkey.as_bytes()).unwrap(),

            pri_key: Rsa::private_key_from_pem(prikey.as_bytes()).unwrap(),

        }

    }



    // public decrypt 

    pub fn public_decrypt(&self, msg:&str) -> String {

        let mut out: [u8; 4096] = [0;4096];

        let decoded = openssl::base64::decode_block(msg).unwrap();

        if let Ok(size) = self.pub_key.public_decrypt(&decoded, &mut out, Padding::PKCS1) {

            let real_size = if size > 4096 {4096} else {size};

            // openssl::base64::encode_block(&out[..real_size])

            String::from_utf8(out[..real_size].to_vec()).unwrap()

        } else {

            String::default()

        }

    }



    // public encrypt 

    pub fn public_encrypt(&self, msg:&str) -> String {

        let mut out: [u8; 4096] = [0;4096];

        if let Ok(size) = self.pub_key.public_encrypt(msg.as_bytes(), &mut out, Padding::PKCS1) {

            let real_size = if size > 4096 {4096}else{size};

            openssl::base64::encode_block(&out[..real_size])

        } else {

            String::default()

        }

    }



    // private encrypt

    pub fn private_encrypt(&self, msg:&str) -> String{

        let mut out: [u8; 4096] = [0;4096];

        if let Ok(size) = self.pri_key.private_encrypt(msg.as_bytes(), &mut out, Padding::PKCS1) {

            let real_size = if size > 4096 {4096}else{size};

            openssl::base64::encode_block(&out[..real_size])

        } else {

            String::default()

        }

    }



    // private decrypt

    pub fn private_decrypt(&self, msg: &str) -> String{

        let mut out: [u8; 4096] = [0;4096];

        let decoded = openssl::base64::decode_block(msg).unwrap();

        if let Ok(size) = self.pri_key.private_decrypt(&decoded, &mut out, Padding::PKCS1) {

            let real_size = if size > 4096 {4096} else {size};

            // openssl::base64::encode_block(&out[..real_size])

            String::from_utf8(out[..real_size].to_vec()).unwrap()

        } else {

            String::default()

        } 

    }



    // sign

    pub fn sign(&self, msg: &str) ->String {

        crypto_utils::hash_hmac(SECRET, msg)

    }

}



#[pymodule]

fn yacrypto(_py: Python, m: &PyModule) -> PyResult<()> {

    m.add_class::<Crypto>()?;

    m.add_function(wrap_pyfunction!(generate_key, m)?).unwrap();

    Ok(())

}





#[cfg(test)]

mod tests {

    use base64;

    #[test]

    fn works(){



        // create rsa

        let rsa = openssl::rsa::Rsa::generate(1024).unwrap();

        // create public key 

        let public_key = rsa.public_key_to_pem().unwrap();

        println!("{:?}", String::from_utf8(public_key.clone()));

        let private_key = rsa.private_key_to_pem().unwrap();





        let data = "hellowo\n\t\rrld";

        // public encrypt 

        let mut buf:Vec<u8> = vec![0;rsa.size() as usize];

        let rsa_pub = openssl::rsa::Rsa::public_key_from_pem(&public_key).unwrap();

        let _ = rsa_pub.public_encrypt(data.as_bytes(), &mut buf , openssl::rsa::Padding::PKCS1);



        // private decrypt => 

        let data = buf;

        let mut buf:Vec<u8> = vec![0;rsa.size() as usize];

        let rsa_pri  = openssl::rsa::Rsa::private_key_from_pem(&private_key).unwrap();

        if let Ok(size) = rsa_pri.private_decrypt(&data, &mut buf, openssl::rsa::Padding::PKCS1){

            let real_size = if size > 1024 {1024} else {size};

            let buf = &buf[..real_size];

            let base64_string = openssl::base64::encode_block(&buf);

            let result = base64::decode(base64_string);

            println!("{:?}",result);

            // println!("buf => {:?}",openssl::base64::encode_block(&buf))



            let echo_str = String::from_utf8((&buf).to_vec()).unwrap();

            println!("{:?}",echo_str);



        }

    }

}

推荐阅读

webpack 从 0 到 1 构建 vue

Ansible 快速入门

  • 0
    感动
  • 0
    路过
  • 0
    高兴
  • 0
    难过
  • 0
    搞笑
  • 0
    无聊
  • 0
    愤怒
  • 0
    同情
热度排行
友情链接