Skip to content

Commit 56c5790

Browse files
jgirardetjgirardet
committed
reformat center, ljsut, rjust args
1 parent 1669d52 commit 56c5790

File tree

3 files changed

+84
-92
lines changed

3 files changed

+84
-92
lines changed

‎tests/snippets/bytes.py

Lines changed: 15 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,8 @@
170170
b"b".center(2, "a")
171171
with assertRaises(TypeError):
172172
b"b".center(2, b"ba")
173+
with assertRaises(TypeError):
174+
b"b".center(b"ba")
173175
assert b"kok".center(5, bytearray(b"x")) == b"xkokx"
174176
b"kok".center(-5)
175177

@@ -200,6 +202,8 @@
200202
b"b".ljust(2, "a")
201203
with assertRaises(TypeError):
202204
b"b".ljust(2, b"ba")
205+
with assertRaises(TypeError):
206+
b"b".ljust(b"ba")
203207
assert b"kok".ljust(5, bytearray(b"x")) == b"kokxx"
204208
assert b"kok".ljust(-5) == b"kok"
205209

@@ -213,23 +217,24 @@
213217
b"||||koki",
214218
b"|||||koki",
215219
]
216-
assert [b"kok".rjust(i, b"|") for i in range(2, 10)] == [
217-
b"kok",
218-
b"kok",
219-
b"|kok",
220-
b"||kok",
221-
b"|||kok",
222-
b"||||kok",
223-
b"|||||kok",
224-
b"||||||kok",
225-
]
220+
assert [b"kok".rjust(i, b"|") for i in range(2, 10)] == [b'kok',
221+
b'kok',
222+
b'|kok',
223+
b'||kok',
224+
b'|||kok',
225+
b'||||kok',
226+
b'|||||kok',
227+
b'||||||kok']
228+
226229

227230

228231
b"kok".rjust(4) == b" kok" # " test no arg"
229232
with assertRaises(TypeError):
230233
b"b".rjust(2, "a")
231234
with assertRaises(TypeError):
232235
b"b".rjust(2, b"ba")
236+
with assertRaises(TypeError):
237+
b"b".rjust(b"ba")
233238
assert b"kok".rjust(5, bytearray(b"x")) == b"xxkok"
234239
assert b"kok".rjust(-5) == b"kok"
235240

‎vm/src/obj/objbyteinner.rs

Lines changed: 60 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,7 @@ impl ByteInnerFindOptions {
116116
elements: &[u8],
117117
vm: &VirtualMachine,
118118
) -> PyResult<(Vec<u8>, Range<usize>)> {
119-
let sub = match try_as_bytes_like(&self.sub.clone()) {
119+
let sub = match try_as_bytes_like(&self.sub) {
120120
Some(value) => value,
121121
None => match_class!(self.sub,
122122
i @ PyInt => vec![i.as_bigint().byte_or(vm)?],
@@ -147,6 +147,56 @@ impl ByteInnerFindOptions {
147147
}
148148
}
149149

150+
#[derive(FromArgs)]
151+
pub struct ByteInnerPaddingOptions {
152+
#[pyarg(positional_only, optional = false)]
153+
width: PyIntRef,
154+
#[pyarg(positional_only, optional = true)]
155+
fillbyte: OptionalArg<PyObjectRef>,
156+
}
157+
158+
impl ByteInnerPaddingOptions {
159+
fn get_value(&self, fn_name: &str, len: usize, vm: &VirtualMachine) -> PyResult<(u8, usize)> {
160+
let fillbyte = if let OptionalArg::Present(v) = &self.fillbyte {
161+
match try_as_byte(&v) {
162+
Some(x) => {
163+
if x.len() == 1 {
164+
x[0]
165+
} else {
166+
return Err(vm.new_type_error(format!(
167+
"{}() argument 2 must be a byte string of length 1, not {}",
168+
fn_name, &v
169+
)));
170+
}
171+
}
172+
None => {
173+
return Err(vm.new_type_error(format!(
174+
"{}() argument 2 must be a byte string of length 1, not {}",
175+
fn_name, &v
176+
)));
177+
}
178+
}
179+
} else {
180+
b' ' // default is space
181+
};
182+
183+
// <0 = no change
184+
let width = if let Some(x) = self.width.as_bigint().to_usize() {
185+
if x <= len {
186+
0
187+
} else {
188+
x
189+
}
190+
} else {
191+
0
192+
};
193+
194+
let diff: usize = if width != 0 { width - len } else { 0 };
195+
196+
Ok((fillbyte, diff))
197+
}
198+
}
199+
150200
impl PyByteInner {
151201
pub fn repr(&self) -> PyResult<String> {
152202
let mut res = String::with_capacity(self.elements.len());
@@ -423,60 +473,13 @@ impl PyByteInner {
423473
.collect::<Vec<u8>>())
424474
}
425475

426-
fn get_center_args(
427-
&self,
428-
width: PyIntRef,
429-
fillbyte: OptionalArg<PyObjectRef>,
430-
fn_name: String,
431-
vm: &VirtualMachine,
432-
) -> PyResult<(u8, usize)> {
433-
let fillbyte = if let OptionalArg::Present(v) = fillbyte {
434-
match try_as_byte(&v) {
435-
Some(x) => {
436-
if x.len() == 1 {
437-
x[0]
438-
} else {
439-
return Err(vm.new_type_error(format!(
440-
"{}() argument 2 must be a byte string of length 1, not {}",
441-
&fn_name, &v
442-
)));
443-
}
444-
}
445-
None => {
446-
return Err(vm.new_type_error(format!(
447-
"{}() argument 2 must be a byte string of length 1, not {}",
448-
&fn_name, &v
449-
)));
450-
}
451-
}
452-
} else {
453-
b' ' // default is space
454-
};
455-
456-
// <0 = no change
457-
let width = if let Some(x) = width.as_bigint().to_usize() {
458-
if x <= self.len() {
459-
0
460-
} else {
461-
x
462-
}
463-
} else {
464-
0
465-
};
466-
467-
let diff: usize = if width != 0 { width - self.len() } else { 0 };
468-
469-
Ok((fillbyte, diff))
470-
}
471-
472476
pub fn center(
473477
&self,
474-
width: PyIntRef,
475-
fillbyte: OptionalArg<PyObjectRef>,
478+
options: ByteInnerPaddingOptions,
476479
vm: &VirtualMachine,
477480
) -> PyResult<Vec<u8>> {
478-
let fn_name = "center".to_string();
479-
let (fillbyte, diff) = self.get_center_args(width, fillbyte, fn_name, vm)?;
481+
// let fn_name = "center".to_string();
482+
let (fillbyte, diff) = options.get_value("center", self.len(), vm)?;
480483

481484
let mut ln: usize = diff / 2;
482485
let mut rn: usize = ln;
@@ -499,12 +502,11 @@ impl PyByteInner {
499502

500503
pub fn ljust(
501504
&self,
502-
width: PyIntRef,
503-
fillbyte: OptionalArg<PyObjectRef>,
505+
options: ByteInnerPaddingOptions,
504506
vm: &VirtualMachine,
505507
) -> PyResult<Vec<u8>> {
506-
let fn_name = "ljust".to_string();
507-
let (fillbyte, diff) = self.get_center_args(width, fillbyte, fn_name, vm)?;
508+
// let fn_name = "ljust".to_string();
509+
let (fillbyte, diff) = options.get_value("ljust", self.len(), vm)?;
508510

509511
// merge all
510512
let mut res = vec![];
@@ -516,12 +518,11 @@ impl PyByteInner {
516518

517519
pub fn rjust(
518520
&self,
519-
width: PyIntRef,
520-
fillbyte: OptionalArg<PyObjectRef>,
521+
options: ByteInnerPaddingOptions,
521522
vm: &VirtualMachine,
522523
) -> PyResult<Vec<u8>> {
523-
let fn_name = "ljust".to_string();
524-
let (fillbyte, diff) = self.get_center_args(width, fillbyte, fn_name, vm)?;
524+
// let fn_name = "rjust".to_string();
525+
let (fillbyte, diff) = options.get_value("rjust", self.len(), vm)?;
525526

526527
// merge all
527528
let mut res = vec![fillbyte; diff];

‎vm/src/obj/objbytes.rs

Lines changed: 9 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
use crate::obj::objint::PyIntRef;
21
use crate::obj::objstr::PyStringRef;
32
use crate::vm::VirtualMachine;
43
use core::cell::Cell;
@@ -7,7 +6,9 @@ use std::ops::Deref;
76
use crate::function::OptionalArg;
87
use crate::pyobject::{PyClassImpl, PyContext, PyIterable, PyObjectRef, PyRef, PyResult, PyValue};
98

10-
use super::objbyteinner::{ByteInnerFindOptions, ByteInnerNewOptions, PyByteInner};
9+
use super::objbyteinner::{
10+
ByteInnerFindOptions, ByteInnerNewOptions, ByteInnerPaddingOptions, PyByteInner,
11+
};
1112
use super::objiter;
1213

1314
use super::objtype::PyClassRef;
@@ -222,33 +223,18 @@ impl PyBytesRef {
222223
}
223224

224225
#[pymethod(name = "center")]
225-
fn center(
226-
self,
227-
width: PyIntRef,
228-
fillbyte: OptionalArg<PyObjectRef>,
229-
vm: &VirtualMachine,
230-
) -> PyResult {
231-
Ok(vm.ctx.new_bytes(self.inner.center(width, fillbyte, vm)?))
226+
fn center(self, options: ByteInnerPaddingOptions, vm: &VirtualMachine) -> PyResult {
227+
Ok(vm.ctx.new_bytes(self.inner.center(options, vm)?))
232228
}
233229

234230
#[pymethod(name = "ljust")]
235-
fn ljust(
236-
self,
237-
width: PyIntRef,
238-
fillbyte: OptionalArg<PyObjectRef>,
239-
vm: &VirtualMachine,
240-
) -> PyResult {
241-
Ok(vm.ctx.new_bytes(self.inner.ljust(width, fillbyte, vm)?))
231+
fn ljust(self, options: ByteInnerPaddingOptions, vm: &VirtualMachine) -> PyResult {
232+
Ok(vm.ctx.new_bytes(self.inner.ljust(options, vm)?))
242233
}
243234

244235
#[pymethod(name = "rjust")]
245-
fn rjust(
246-
self,
247-
width: PyIntRef,
248-
fillbyte: OptionalArg<PyObjectRef>,
249-
vm: &VirtualMachine,
250-
) -> PyResult {
251-
Ok(vm.ctx.new_bytes(self.inner.rjust(width, fillbyte, vm)?))
236+
fn rjust(self, options: ByteInnerPaddingOptions, vm: &VirtualMachine) -> PyResult {
237+
Ok(vm.ctx.new_bytes(self.inner.rjust(options, vm)?))
252238
}
253239

254240
#[pymethod(name = "count")]

0 commit comments

Comments
 (0)