|
2 | 2 | """ |
3 | 3 |
|
4 | 4 | import argparse |
5 | | - |
| 5 | +import os |
6 | 6 | from html2image import Html2Image |
7 | 7 |
|
8 | 8 |
|
9 | | -def main(): |
10 | | - |
11 | | - def size_type(string): |
12 | | - try: |
13 | | - x, y = map(int, string.split(',')) |
14 | | - return x, y |
15 | | - except Exception: |
| 9 | +def size_type(string): |
| 10 | + try: |
| 11 | + width, height = map(int, string.split(',')) |
| 12 | + if width <= 0 or height <= 0: |
16 | 13 | raise argparse.ArgumentTypeError( |
17 | | - f"size should be int,int, instead got {string}" |
| 14 | + 'Width and height must be positive integers.' |
18 | 15 | ) |
| 16 | + return width, height |
| 17 | + except ValueError: # incorrect number of values |
| 18 | + raise argparse.ArgumentTypeError( |
| 19 | + f"Size should be W,H (e.g., 1920,1080), instead got '{string}'" |
| 20 | + ) |
| 21 | + except Exception as e: # unexpected errors |
| 22 | + raise argparse.ArgumentTypeError(f"Invalid size format '{string}': {e}") |
19 | 23 |
|
20 | | - parser = argparse.ArgumentParser() |
21 | 24 |
|
22 | | - parser.add_argument('-U', '--url', nargs='*', required=False, default=[]) |
23 | | - parser.add_argument('-H', '--html', nargs='*', required=False, default=[]) |
24 | | - parser.add_argument('-C', '--css', nargs='*', required=False, default=[]) |
25 | | - parser.add_argument('-O', '--other', nargs='*', required=False, default=[]) |
| 25 | +def main(): |
| 26 | + parser = argparse.ArgumentParser( |
| 27 | + description='Generate images from HTML/CSS or URLs using the html2image library.', |
| 28 | + formatter_class=argparse.ArgumentDefaultsHelpFormatter |
| 29 | + ) |
| 30 | + |
| 31 | + # Html2Image instantiation arguments |
| 32 | + group_hti_init = parser.add_argument_group('Html2Image Instance Configuration') |
| 33 | + group_hti_init.add_argument( |
| 34 | + '--output-path', '-o', |
| 35 | + default=os.getcwd(), |
| 36 | + help='Directory to save screenshots.' |
| 37 | + ) |
26 | 38 |
|
27 | | - parser.add_argument( |
28 | | - '-S', '--save-as', nargs='*', required=False, default="screenshot.png" |
| 39 | + # TODO : this list is duplicated from browser_map in html2image.py |
| 40 | + browser_choices = [ |
| 41 | + 'chrome', 'chromium', 'google-chrome', 'google-chrome-stable', |
| 42 | + 'googlechrome', 'edge', 'chrome-cdp', 'chromium-cdp' |
| 43 | + ] |
| 44 | + group_hti_init.add_argument( |
| 45 | + '--browser', |
| 46 | + default='chrome', |
| 47 | + choices=browser_choices, |
| 48 | + help='Browser to use for screenshots.' |
29 | 49 | ) |
30 | | - parser.add_argument( |
31 | | - '-s', '--size', nargs='*', required=False, default=[], type=size_type |
| 50 | + group_hti_init.add_argument( |
| 51 | + '--browser-executable', |
| 52 | + default=None, |
| 53 | + help='Path to the browser executable. Auto-detected if not provided.' |
| 54 | + ) |
| 55 | + group_hti_init.add_argument( |
| 56 | + '--cdp-port', |
| 57 | + type=int, |
| 58 | + default=None, |
| 59 | + help='CDP port for CDP-enabled browsers (e.g., chrome-cdp). Default is library-dependent.' |
| 60 | + ) |
| 61 | + group_hti_init.add_argument( |
| 62 | + '--temp-path', |
| 63 | + default=None, |
| 64 | + help="Directory for temporary files. Defaults to system temp directory within an 'html2image' subfolder." |
| 65 | + ) |
| 66 | + group_hti_init.add_argument( |
| 67 | + '--keep-temp-files', |
| 68 | + action='store_true', |
| 69 | + help='Do not delete temporary files after screenshot generation.' |
| 70 | + ) |
| 71 | + group_hti_init.add_argument( |
| 72 | + '--custom-flags', |
| 73 | + nargs='*', |
| 74 | + default=[], # If not provided, defaults are used |
| 75 | + help="Custom flags to pass to the browser (e.g., '--no-sandbox' '--disable-gpu'). If provided, these flags will be used." |
32 | 76 | ) |
33 | 77 |
|
34 | | - parser.add_argument('-o', '--output_path', required=False) |
| 78 | + # Screenshot sources arguments |
| 79 | + group_sources = parser.add_argument_group('Screenshot Sources (at least one type is required)') |
| 80 | + group_sources.add_argument( |
| 81 | + '--url', '-U', |
| 82 | + nargs='*', default=[], |
| 83 | + metavar='URL', |
| 84 | + help='URL(s) to screenshot.' |
| 85 | + ) |
| 86 | + group_sources.add_argument( |
| 87 | + '--html-file', |
| 88 | + nargs='*', default=[], |
| 89 | + metavar='FILE', |
| 90 | + help='HTML file(s) to screenshot.' |
| 91 | + ) |
| 92 | + group_sources.add_argument( |
| 93 | + '--html-string', |
| 94 | + nargs='*', default=[], |
| 95 | + metavar='STRING', |
| 96 | + help='HTML string(s) to screenshot.' |
| 97 | + ) |
| 98 | + group_sources.add_argument( |
| 99 | + '--css-file', |
| 100 | + nargs='*', default=[], |
| 101 | + metavar='FILE', |
| 102 | + help='CSS file(s) to load. Used by HTML files or applied with HTML strings.' |
| 103 | + ) |
| 104 | + group_sources.add_argument( |
| 105 | + '--css-string', |
| 106 | + nargs='*', default=[], |
| 107 | + metavar='STRING', |
| 108 | + help='CSS string(s) to apply. Combined and used with HTML strings.' |
| 109 | + ) |
| 110 | + group_sources.add_argument( |
| 111 | + '--other-file', '-O', |
| 112 | + nargs='*', default=[], |
| 113 | + metavar='FILE', |
| 114 | + help='Other file(s) to screenshot (e.g., SVG).' |
| 115 | + ) |
35 | 116 |
|
36 | | - parser.add_argument('-q', '--quiet', required=False, action="store_true") |
37 | | - parser.add_argument('-v', '--verbose', required=False, action="store_true") |
| 117 | + # Screenshot output control arguments |
| 118 | + group_output_ctrl = parser.add_argument_group('Screenshot Output Options') |
| 119 | + group_output_ctrl.add_argument( |
| 120 | + '--save-as', '-S', |
| 121 | + nargs='*', default=None, # html2image handles default naming if only one source |
| 122 | + metavar='FILENAME', |
| 123 | + help='Filename(s) for the output images. If not provided or fewer names than items, names are auto-generated.' |
| 124 | + ) |
| 125 | + group_output_ctrl.add_argument( |
| 126 | + '--size', '-s', |
| 127 | + nargs='*', default=[], |
| 128 | + type=size_type, |
| 129 | + metavar='W,H', |
| 130 | + help="Size(s) for screenshots as W,H. If one W,H pair is given, it applies to all. If multiple, they apply to corresponding screenshots; if fewer pairs than items, the last is repeated. If omitted, (1920,1080) is used." |
| 131 | + ) |
38 | 132 |
|
39 | | - # parser.add_argument('--browser', required=False) |
40 | | - parser.add_argument('--chrome_path', required=False) |
41 | | - # parser.add_argument('--firefox_path', required=False) |
42 | | - parser.add_argument('--temp_path', required=False) |
43 | | - parser.add_argument('--custom_flags', required=False) |
| 133 | + # General arguments |
| 134 | + group_general = parser.add_argument_group('General Options') |
| 135 | + group_general.add_argument( |
| 136 | + '--quiet', '-q', |
| 137 | + action='store_true', |
| 138 | + help='Suppress output from browsers (sets disable_logging=True).' |
| 139 | + ) |
| 140 | + group_general.add_argument( |
| 141 | + '--verbose', '-v', |
| 142 | + action='store_true', |
| 143 | + help='Enable verbose output, including browser commands if supported by the browser handler.' |
| 144 | + ) |
44 | 145 |
|
45 | 146 | args = parser.parse_args() |
46 | 147 |
|
| 148 | + # Prepare Html2Image() |
| 149 | + hti_kwargs = { |
| 150 | + 'output_path': args.output_path, |
| 151 | + 'browser': args.browser, |
| 152 | + 'browser_executable': args.browser_executable, |
| 153 | + 'custom_flags': [cf.replace("'", '') for cf in args.custom_flags], |
| 154 | + 'disable_logging': args.quiet, |
| 155 | + 'temp_path': args.temp_path, |
| 156 | + 'keep_temp_files': args.keep_temp_files, |
| 157 | + } |
| 158 | + |
| 159 | + # Only pass cdp_port if a CDP browser is likely selected and port is given |
| 160 | + if args.cdp_port and 'cdp' in args.browser.lower(): |
| 161 | + hti_kwargs['browser_cdp_port'] = args.cdp_port |
| 162 | + elif args.cdp_port: |
| 163 | + print( |
| 164 | + f"Warning: --cdp-port ({args.cdp_port}) was specified, but the selected browser ('{args.browser}') might not be a CDP browser." |
| 165 | + ) |
| 166 | + |
47 | 167 | try: |
48 | | - hti = Html2Image(disable_logging=args.quiet) |
| 168 | + # Filter out None values so defaults are used for those specific kwargs |
| 169 | + # keep_temp_files and disable_logging are bools, always pass them. |
| 170 | + # custom_flags should be passed even if None, so Html2Image can use its defaults or an empty list. |
| 171 | + active_hti_kwargs = { |
| 172 | + k: v for k, v in hti_kwargs.items() |
| 173 | + if v is not None or k in ['keep_temp_files', 'disable_logging', 'custom_flags'] |
| 174 | + } |
| 175 | + hti = Html2Image(**active_hti_kwargs) |
| 176 | + |
49 | 177 | except Exception as e: |
50 | | - print('Could not instanciate html2image.') |
51 | | - print(e) |
| 178 | + print(f'Error: Could not instantiate Html2Image: {e}') |
52 | 179 | exit(1) |
53 | 180 |
|
54 | 181 | if args.verbose: |
55 | | - print(f'args = {args}') |
56 | | - hti.browser.print_command = True |
57 | | - |
58 | | - if args.output_path: |
59 | | - hti.output_path = args.output_path |
| 182 | + # The `print_command` attribute is specific to ChromiumHeadless. |
| 183 | + # CDP browsers print logs internally. |
| 184 | + if hasattr(hti.browser, 'print_command'): |
| 185 | + hti.browser.print_command = True |
| 186 | + print('Verbose mode: Browser commands will be printed for compatible handlers.') |
| 187 | + else: |
| 188 | + print('Verbose mode enabled. Note: Detailed browser command printing depends on the selected browser handler.') |
| 189 | + |
| 190 | + has_sources = any([ |
| 191 | + args.url, args.html_file, args.html_string, args.other_file |
| 192 | + ]) |
| 193 | + |
| 194 | + # Print help message if no sources were passed |
| 195 | + if not has_sources: |
| 196 | + print('Error: No screenshot sources (URL, HTML file/string, other file) provided.') |
| 197 | + parser.print_usage() |
| 198 | + exit(1) |
60 | 199 |
|
61 | | - if args.chrome_path: |
62 | | - hti.chrome_path = args.chrome_path |
| 200 | + # Perform screenshot |
| 201 | + screenshot_kwargs = { |
| 202 | + 'url': args.url, |
| 203 | + 'html_file': args.html_file, |
| 204 | + 'html_str': args.html_string, |
| 205 | + 'css_file': args.css_file, |
| 206 | + 'css_str': args.css_string, |
| 207 | + 'other_file': args.other_file, |
| 208 | + 'size': args.size, # Pass the list of sizes directly from the --size CLI arg |
| 209 | + } |
63 | 210 |
|
64 | | - if args.custom_flags: |
65 | | - hti.browser.flags = args.custom_flags |
| 211 | + if args.save_as is not None: |
| 212 | + screenshot_kwargs['save_as'] = args.save_as |
66 | 213 |
|
67 | | - if args.temp_path: |
68 | | - hti.temp_path = args.temp_path |
| 214 | + try: |
| 215 | + if args.verbose: |
| 216 | + print('--- Html2Image Instance Configuration ---') |
| 217 | + for k, v in active_hti_kwargs.items(): |
| 218 | + print(f' {k}: {v}') |
| 219 | + print('--- Screenshot Call Arguments ---') |
| 220 | + for k, v in screenshot_kwargs.items(): |
| 221 | + if v or k == 'size': # print if list not empty, or always for size |
| 222 | + print(f' {k}: {v}') |
| 223 | + |
| 224 | + paths = hti.screenshot(**screenshot_kwargs) |
| 225 | + |
| 226 | + if not args.quiet: |
| 227 | + print(f'Successfully created {len(paths)} image(s):') |
| 228 | + for path in paths: |
| 229 | + print(f' {path}') |
| 230 | + |
| 231 | + except FileNotFoundError as e: |
| 232 | + print(f'Error: A required file was not found: {e}') |
| 233 | + exit(1) |
69 | 234 |
|
70 | | - paths = hti.screenshot( |
71 | | - html_file=args.html, css_file=args.css, other_file=args.other, |
72 | | - url=args.url, save_as=args.save_as, size=args.size |
73 | | - ) |
| 235 | + except ValueError as e: # Can be raised by browser screenshot method for bad size etc. |
| 236 | + print(f'Error: Invalid value encountered: {e}') |
| 237 | + exit(1) |
74 | 238 |
|
75 | | - if not args.quiet: |
76 | | - print(f'Created {len(paths)} file(s):') |
77 | | - for path in paths: |
78 | | - print(f'\t{path}') |
| 239 | + except Exception as e: |
| 240 | + print(f'An unexpected error occurred during screenshotting: {e}') |
| 241 | + if args.verbose: |
| 242 | + import traceback |
| 243 | + traceback.print_exc() |
| 244 | + exit(1) |
79 | 245 |
|
80 | 246 |
|
81 | | -if __name__ == "__main__": |
| 247 | +if __name__ == '__main__': |
82 | 248 | main() |
0 commit comments