@@ -127,6 +127,62 @@ const CLEANUP_SELECTORS = {
127127 ] ,
128128} ;
129129
130+ export interface PrettyScreenshotOptions {
131+ scrollTo ?: string ;
132+ doCleanup : boolean ;
133+ hideSelectors : string [ ] ;
134+ viewportWidth ?: number ;
135+ outputPath ?: string ;
136+ }
137+
138+ function isLikelyScreenshotOutputPath ( arg : string ) : boolean {
139+ return arg . startsWith ( '/' )
140+ || arg . startsWith ( './' )
141+ || arg . startsWith ( '../' )
142+ || arg . startsWith ( '~' )
143+ || / \. ( p n g | j p e ? g | w e b p ) $ / i. test ( arg ) ;
144+ }
145+
146+ export function parsePrettyScreenshotArgs ( args : string [ ] ) : PrettyScreenshotOptions {
147+ const options : PrettyScreenshotOptions = {
148+ doCleanup : false ,
149+ hideSelectors : [ ] ,
150+ } ;
151+
152+ for ( let i = 0 ; i < args . length ; i ++ ) {
153+ if ( args [ i ] === '--scroll-to' && i + 1 < args . length ) {
154+ options . scrollTo = args [ ++ i ] ;
155+ } else if ( args [ i ] === '--cleanup' ) {
156+ options . doCleanup = true ;
157+ } else if ( args [ i ] === '--hide' && i + 1 < args . length ) {
158+ const values : string [ ] = [ ] ;
159+ i ++ ;
160+ while ( i < args . length && ! args [ i ] . startsWith ( '--' ) ) {
161+ values . push ( args [ i ] ) ;
162+ i ++ ;
163+ }
164+ i -- ; // Back up since the for loop will increment.
165+
166+ if ( ! options . outputPath && values . length > 1 ) {
167+ const last = values [ values . length - 1 ] ;
168+ if ( isLikelyScreenshotOutputPath ( last ) ) {
169+ options . outputPath = values . pop ( ) ;
170+ }
171+ }
172+ options . hideSelectors . push ( ...values ) ;
173+ } else if ( args [ i ] === '--width' && i + 1 < args . length ) {
174+ options . viewportWidth = parseInt ( args [ ++ i ] , 10 ) ;
175+ if ( isNaN ( options . viewportWidth ) ) throw new Error ( '--width must be a number' ) ;
176+ } else if ( ! args [ i ] . startsWith ( '--' ) ) {
177+ options . outputPath = args [ i ] ;
178+ } else {
179+ throw new Error ( `Unknown prettyscreenshot flag: ${ args [ i ] } ` ) ;
180+ }
181+ }
182+
183+ return options ;
184+ }
185+
130186export async function handleWriteCommand (
131187 command : string ,
132188 args : string [ ] ,
@@ -992,35 +1048,9 @@ export async function handleWriteCommand(
9921048 }
9931049
9941050 case 'prettyscreenshot' : {
995- // Parse flags
996- let scrollTo : string | undefined ;
997- let doCleanup = false ;
998- const hideSelectors : string [ ] = [ ] ;
999- let viewportWidth : number | undefined ;
1000- let outputPath : string | undefined ;
1001-
1002- for ( let i = 0 ; i < args . length ; i ++ ) {
1003- if ( args [ i ] === '--scroll-to' && i + 1 < args . length ) {
1004- scrollTo = args [ ++ i ] ;
1005- } else if ( args [ i ] === '--cleanup' ) {
1006- doCleanup = true ;
1007- } else if ( args [ i ] === '--hide' && i + 1 < args . length ) {
1008- // Collect all following non-flag args as selectors to hide
1009- i ++ ;
1010- while ( i < args . length && ! args [ i ] . startsWith ( '--' ) ) {
1011- hideSelectors . push ( args [ i ] ) ;
1012- i ++ ;
1013- }
1014- i -- ; // Back up since the for loop will increment
1015- } else if ( args [ i ] === '--width' && i + 1 < args . length ) {
1016- viewportWidth = parseInt ( args [ ++ i ] , 10 ) ;
1017- if ( isNaN ( viewportWidth ) ) throw new Error ( '--width must be a number' ) ;
1018- } else if ( ! args [ i ] . startsWith ( '--' ) ) {
1019- outputPath = args [ i ] ;
1020- } else {
1021- throw new Error ( `Unknown prettyscreenshot flag: ${ args [ i ] } ` ) ;
1022- }
1023- }
1051+ const prettyOptions = parsePrettyScreenshotArgs ( args ) ;
1052+ const { scrollTo, doCleanup, hideSelectors, viewportWidth } = prettyOptions ;
1053+ let { outputPath } = prettyOptions ;
10241054
10251055 // Default output path
10261056 if ( ! outputPath ) {
0 commit comments