rust 표준 라이브러리에 사용자 정의 메소드 추가하는 법

가령 std::TcpStream 같은 표준 라이브러리에 새로운 메소드를 추가하고 싶다면, 트레잇을 새로 정의하여 확장할 수 있다.

메소드 추가

use std::io::{self, Write};
use std::net::TcpStream;

pub trait TcpStreamExt {
    /// 버퍼를 모두 쓰거나 에러가 발생할 때까지 쓰는 메소드
    fn write_n(&mut self, buf: &[u8]) -> io::Result<usize>;
}

impl TcpStreamExt for TcpStream {
    fn write_n(&mut self, mut buf: &[u8]) -> io::Result<usize> {
        let mut written_n = 0;
        while !buf.is_empty() {
            match self.write(buf) {
                Ok(0) => return Ok(0), // 아무 것도 쓰지 못한 경우 종료
                Ok(n) => {
                    buf = &buf[n..];
                    written_n += n;
                }
                Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {
                    // Interrupted 에러 발생 시 재시도
                    continue;
                }
                Err(e) => return Err(e),
            }
        }
        Ok(written_n)
    }
}

이렇게 하면 표준 라이브러리 변경 없이 원하는 메소드를 안전하게 추가할 수 있다. 또한 특정 프로젝트 내에서만 확장 메소드를 정의할 수 있어 필요하지 않은 경우를 최소화 할 수 있다.

사용법

fn main() -> std::io::Result<()> {
    use std::net::TcpStream;

    let mut stream = TcpStream::connect("example.com:80")?;
    let data = b"Hello, world!";

    // write_n 메소드를 사용
    let bytes_written = stream.write_n(data)?;
    println!("Bytes written: {}", bytes_written);

    Ok(())
}