Skip to content

Commit 41e5c89

Browse files
dsherretdsherret
committed
fix: allow negative numbers in head -n 
1 parent 0fde85b commit 41e5c89

File tree

1 file changed

+240
-40
lines changed

1 file changed

+240
-40
lines changed

‎src/shell/commands/head.rs‎

Lines changed: 240 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -81,47 +81,147 @@ fn copy_lines<F: FnMut(&mut [u8]) -> Result<usize>>(
8181
Ok(ExecuteResult::from_exit_code(0))
8282
}
8383

84+
fn copy_all_but_last_lines<F: FnMut(&mut [u8]) -> Result<usize>>(
85+
writer: &mut ShellPipeWriter,
86+
skip_last: u64,
87+
kill_signal: &KillSignal,
88+
mut read: F,
89+
) -> Result<ExecuteResult> {
90+
// read all content first
91+
let mut content = Vec::new();
92+
let mut buffer = vec![0; 512];
93+
loop {
94+
if let Some(exit_code) = kill_signal.aborted_code() {
95+
return Ok(ExecuteResult::from_exit_code(exit_code));
96+
}
97+
let read_bytes = read(&mut buffer)?;
98+
if read_bytes == 0 {
99+
break;
100+
}
101+
content.extend_from_slice(&buffer[..read_bytes]);
102+
}
103+
104+
// count total lines
105+
let total_lines = content.iter().filter(|&&b| b == b'\n').count() as u64;
106+
107+
// output all but the last N lines
108+
if total_lines <= skip_last {
109+
return Ok(ExecuteResult::from_exit_code(0));
110+
}
111+
let lines_to_print = total_lines - skip_last;
112+
113+
let mut line_count = 0u64;
114+
let mut start = 0;
115+
for (i, &b) in content.iter().enumerate() {
116+
if let Some(exit_code) = kill_signal.aborted_code() {
117+
return Ok(ExecuteResult::from_exit_code(exit_code));
118+
}
119+
if b == b'\n' {
120+
line_count += 1;
121+
if line_count <= lines_to_print {
122+
writer.write_all(&content[start..=i])?;
123+
}
124+
start = i + 1;
125+
if line_count >= lines_to_print {
126+
break;
127+
}
128+
}
129+
}
130+
131+
Ok(ExecuteResult::from_exit_code(0))
132+
}
133+
84134
fn execute_head(mut context: ShellCommandContext) -> Result<ExecuteResult> {
85135
let flags = parse_args(&context.args)?;
86-
if flags.path == "-" {
87-
copy_lines(
88-
&mut context.stdout,
89-
flags.lines,
90-
context.state.kill_signal(),
91-
|buf| context.stdin.read(buf),
92-
512,
93-
)
94-
} else {
95-
let path = flags.path;
96-
match File::open(context.state.cwd().join(path)) {
97-
Ok(mut file) => copy_lines(
98-
&mut context.stdout,
99-
flags.lines,
100-
context.state.kill_signal(),
101-
|buf| file.read(buf).map_err(Into::into),
102-
512,
103-
),
104-
Err(err) => {
105-
context.stderr.write_line(&format!(
106-
"head: {}: {}",
107-
path.to_string_lossy(),
108-
err
109-
))?;
110-
Ok(ExecuteResult::from_exit_code(1))
136+
match flags.lines {
137+
LineCount::First(max_lines) => {
138+
if flags.path == "-" {
139+
copy_lines(
140+
&mut context.stdout,
141+
max_lines,
142+
context.state.kill_signal(),
143+
|buf| context.stdin.read(buf),
144+
512,
145+
)
146+
} else {
147+
let path = flags.path;
148+
match File::open(context.state.cwd().join(path)) {
149+
Ok(mut file) => copy_lines(
150+
&mut context.stdout,
151+
max_lines,
152+
context.state.kill_signal(),
153+
|buf| file.read(buf).map_err(Into::into),
154+
512,
155+
),
156+
Err(err) => {
157+
context.stderr.write_line(&format!(
158+
"head: {}: {}",
159+
path.to_string_lossy(),
160+
err
161+
))?;
162+
Ok(ExecuteResult::from_exit_code(1))
163+
}
164+
}
165+
}
166+
}
167+
LineCount::AllButLast(skip_last) => {
168+
if flags.path == "-" {
169+
copy_all_but_last_lines(
170+
&mut context.stdout,
171+
skip_last,
172+
context.state.kill_signal(),
173+
|buf| context.stdin.read(buf),
174+
)
175+
} else {
176+
let path = flags.path;
177+
match File::open(context.state.cwd().join(path)) {
178+
Ok(mut file) => copy_all_but_last_lines(
179+
&mut context.stdout,
180+
skip_last,
181+
context.state.kill_signal(),
182+
|buf| file.read(buf).map_err(Into::into),
183+
),
184+
Err(err) => {
185+
context.stderr.write_line(&format!(
186+
"head: {}: {}",
187+
path.to_string_lossy(),
188+
err
189+
))?;
190+
Ok(ExecuteResult::from_exit_code(1))
191+
}
192+
}
111193
}
112194
}
113195
}
114196
}
115197

198+
#[derive(Debug, PartialEq, Clone, Copy)]
199+
enum LineCount {
200+
/// print first N lines
201+
First(u64),
202+
/// print all but last N lines
203+
AllButLast(u64),
204+
}
205+
116206
#[derive(Debug, PartialEq)]
117207
struct HeadFlags<'a> {
118208
path: &'a OsStr,
119-
lines: u64,
209+
lines: LineCount,
210+
}
211+
212+
fn parse_line_count(s: &str) -> Result<LineCount> {
213+
if let Some(rest) = s.strip_prefix('-') {
214+
let num = rest.parse::<u64>()?;
215+
Ok(LineCount::AllButLast(num))
216+
} else {
217+
let num = s.parse::<u64>()?;
218+
Ok(LineCount::First(num))
219+
}
120220
}
121221

122222
fn parse_args<'a>(args: &'a [OsString]) -> Result<HeadFlags<'a>> {
123223
let mut path: Option<&'a OsStr> = None;
124-
let mut lines: Option<u64> = None;
224+
let mut lines: Option<LineCount> = None;
125225
let mut iterator = parse_arg_kinds(args).into_iter();
126226
while let Some(arg) = iterator.next() {
127227
match arg {
@@ -137,9 +237,11 @@ fn parse_args<'a>(args: &'a [OsString]) -> Result<HeadFlags<'a>> {
137237
}
138238
ArgKind::ShortFlag('n') => match iterator.next() {
139239
Some(ArgKind::Arg(arg)) => {
140-
let num = arg.to_str().and_then(|a| a.parse::<u64>().ok());
141-
if let Some(num) = num {
142-
lines = Some(num);
240+
if let Some(s) = arg.to_str() {
241+
match parse_line_count(s) {
242+
Ok(count) => lines = Some(count),
243+
Err(_) => bail!("expected a numeric value following -n"),
244+
}
143245
} else {
144246
bail!("expected a numeric value following -n")
145247
}
@@ -150,7 +252,7 @@ fn parse_args<'a>(args: &'a [OsString]) -> Result<HeadFlags<'a>> {
150252
if flag == "lines" || flag == "lines=" {
151253
bail!("expected a value for --lines");
152254
} else if let Some(arg) = flag.strip_prefix("lines=") {
153-
lines = Some(arg.parse::<u64>()?);
255+
lines = Some(parse_line_count(arg)?);
154256
} else {
155257
arg.bail_unsupported()?
156258
}
@@ -161,7 +263,7 @@ fn parse_args<'a>(args: &'a [OsString]) -> Result<HeadFlags<'a>> {
161263

162264
Ok(HeadFlags {
163265
path: path.unwrap_or(OsStr::new("-")),
164-
lines: lines.unwrap_or(10),
266+
lines: lines.unwrap_or(LineCount::First(10)),
165267
})
166268
}
167269

@@ -231,56 +333,78 @@ mod test {
231333
parse_args(&[]).unwrap(),
232334
HeadFlags {
233335
path: OsStr::new("-"),
234-
lines: 10
336+
lines: LineCount::First(10)
235337
}
236338
);
237339
assert_eq!(
238340
parse_args(&["-n".into(), "5".into()]).unwrap(),
239341
HeadFlags {
240342
path: OsStr::new("-"),
241-
lines: 5
343+
lines: LineCount::First(5)
242344
}
243345
);
244346
assert_eq!(
245347
parse_args(&["--lines=5".into()]).unwrap(),
246348
HeadFlags {
247349
path: OsStr::new("-"),
248-
lines: 5
350+
lines: LineCount::First(5)
249351
}
250352
);
251353
assert_eq!(
252354
parse_args(&["path".into()]).unwrap(),
253355
HeadFlags {
254356
path: OsStr::new("path"),
255-
lines: 10
357+
lines: LineCount::First(10)
256358
}
257359
);
258360
assert_eq!(
259361
parse_args(&["-n".into(), "5".into(), "path".into()]).unwrap(),
260362
HeadFlags {
261363
path: OsStr::new("path"),
262-
lines: 5
364+
lines: LineCount::First(5)
263365
}
264366
);
265367
assert_eq!(
266368
parse_args(&["--lines=5".into(), "path".into()]).unwrap(),
267369
HeadFlags {
268370
path: OsStr::new("path"),
269-
lines: 5
371+
lines: LineCount::First(5)
270372
}
271373
);
272374
assert_eq!(
273375
parse_args(&["path".into(), "-n".into(), "5".into()]).unwrap(),
274376
HeadFlags {
275377
path: OsStr::new("path"),
276-
lines: 5
378+
lines: LineCount::First(5)
277379
}
278380
);
279381
assert_eq!(
280382
parse_args(&["path".into(), "--lines=5".into()]).unwrap(),
281383
HeadFlags {
282384
path: OsStr::new("path"),
283-
lines: 5
385+
lines: LineCount::First(5)
386+
}
387+
);
388+
// negative line counts (all but last N)
389+
assert_eq!(
390+
parse_args(&["-n".into(), "-1".into()]).unwrap(),
391+
HeadFlags {
392+
path: OsStr::new("-"),
393+
lines: LineCount::AllButLast(1)
394+
}
395+
);
396+
assert_eq!(
397+
parse_args(&["-n".into(), "-5".into(), "path".into()]).unwrap(),
398+
HeadFlags {
399+
path: OsStr::new("path"),
400+
lines: LineCount::AllButLast(5)
401+
}
402+
);
403+
assert_eq!(
404+
parse_args(&["--lines=-3".into()]).unwrap(),
405+
HeadFlags {
406+
path: OsStr::new("-"),
407+
lines: LineCount::AllButLast(3)
284408
}
285409
);
286410
assert_eq!(
@@ -304,4 +428,80 @@ mod test {
304428
"unsupported flag: -t"
305429
);
306430
}
431+
432+
#[tokio::test]
433+
async fn copies_all_but_last_lines() {
434+
let (reader, mut writer) = pipe();
435+
let reader_handle = reader.pipe_to_string_handle();
436+
let data = b"line1\nline2\nline3\nline4\nline5\n";
437+
let mut offset = 0;
438+
let result = copy_all_but_last_lines(
439+
&mut writer,
440+
1,
441+
&KillSignal::default(),
442+
|buffer| {
443+
if offset >= data.len() {
444+
return Ok(0);
445+
}
446+
let read_length = min(buffer.len(), data.len() - offset);
447+
buffer[..read_length].copy_from_slice(&data[offset..(offset + read_length)]);
448+
offset += read_length;
449+
Ok(read_length)
450+
},
451+
);
452+
drop(writer);
453+
assert_eq!(reader_handle.await.unwrap(), "line1\nline2\nline3\nline4\n");
454+
assert_eq!(result.unwrap().into_exit_code_and_handles().0, 0);
455+
}
456+
457+
#[tokio::test]
458+
async fn copies_all_but_last_two_lines() {
459+
let (reader, mut writer) = pipe();
460+
let reader_handle = reader.pipe_to_string_handle();
461+
let data = b"line1\nline2\nline3\nline4\nline5\n";
462+
let mut offset = 0;
463+
let result = copy_all_but_last_lines(
464+
&mut writer,
465+
2,
466+
&KillSignal::default(),
467+
|buffer| {
468+
if offset >= data.len() {
469+
return Ok(0);
470+
}
471+
let read_length = min(buffer.len(), data.len() - offset);
472+
buffer[..read_length].copy_from_slice(&data[offset..(offset + read_length)]);
473+
offset += read_length;
474+
Ok(read_length)
475+
},
476+
);
477+
drop(writer);
478+
assert_eq!(reader_handle.await.unwrap(), "line1\nline2\nline3\n");
479+
assert_eq!(result.unwrap().into_exit_code_and_handles().0, 0);
480+
}
481+
482+
#[tokio::test]
483+
async fn copies_all_but_last_lines_when_skip_exceeds_total() {
484+
let (reader, mut writer) = pipe();
485+
let reader_handle = reader.pipe_to_string_handle();
486+
let data = b"line1\nline2\n";
487+
let mut offset = 0;
488+
let result = copy_all_but_last_lines(
489+
&mut writer,
490+
5,
491+
&KillSignal::default(),
492+
|buffer| {
493+
if offset >= data.len() {
494+
return Ok(0);
495+
}
496+
let read_length = min(buffer.len(), data.len() - offset);
497+
buffer[..read_length].copy_from_slice(&data[offset..(offset + read_length)]);
498+
offset += read_length;
499+
Ok(read_length)
500+
},
501+
);
502+
drop(writer);
503+
// when skip_last >= total_lines, output should be empty
504+
assert_eq!(reader_handle.await.unwrap(), "");
505+
assert_eq!(result.unwrap().into_exit_code_and_handles().0, 0);
506+
}
307507
}

0 commit comments

Comments
 (0)