Fix argument parsing to support ml images as well

[Deependra-Patel]

Also migrated a unit test to Python3

TESTED=Unit test is now passing. Created a custom image with ml base image.

[gemini-code-assist]

Summary of Changes

Hello @Deependra-Patel, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request enhances the argument parsing utility to support machine learning (ML) specific image versions, such as 'ml-ubuntu'. It refines the regular expression patterns used for image version identification and improves the string manipulation logic for image path inference. Additionally, the associated unit tests have been updated to align with Python 3 standards and to validate the new ML image parsing capabilities, ensuring the robustness of the argument parsing system.

Highlights

• ML Image Version Support: Enabled parsing for ML-specific image versions by updating regular expressions to correctly identify 'ml' prefixes in image names.
• Argument Parsing Logic Refinement: Improved argument parsing logic by ensuring that version string replacements only occur on the first instance, preventing unintended modifications.
• Unit Test Modernization: Modernized the unit test suite by migrating it to Python 3's 'argparse' module and updating test assertions to directly compare 'argparse.Namespace' objects, enhancing test reliability and readability.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• custom_image_utils/args_inferer.py
  • Modified the 'replace' method to limit replacements to the first occurrence of the pattern '-\d+-' when constructing the filter argument for Dataproc image paths.
• custom_image_utils/args_parser.py
  • Updated '_VERSION_REGEX' to allow for 'ml' prefixes (e.g., '-ml-') in version strings.
  • Modified '_LATEST_FROM_MINOR_VERSION' regex to correctly parse minor versions that include optional 'ml' prefixes before the OS distribution.
• tests/test_args_parser.py
  • Replaced the import of the deprecated 'exceptions' module with 'argparse'.
  • Refactored the '_make_expected_result' helper function to directly return an 'argparse.Namespace' object, simplifying test result construction.
  • Updated various test assertions to compare 'argparse.Namespace' objects directly instead of their string representations.
  • Added a new test case to validate the parsing of 'ml' prefixed Dataproc versions (e.g., '2.3-ml-ubuntu22').

Activity

• Unit tests were updated and are now passing.
• A custom image with an ML base image was successfully created, confirming the fix.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces support for ML image names in argument parsing, which is a valuable addition. The changes to the regular expressions in args_parser.py and the logic in args_inferer.py are correct and well-implemented to handle the new image name formats.

A significant part of this PR is the migration of tests/test_args_parser.py to be Python 3 compatible and the substantial refactoring of the tests. Moving from string-based comparisons of Namespace objects to direct object comparison is a major improvement in test quality, making them more robust and readable. The cleanup of the test helper _make_expected_result is particularly noteworthy.

I have one suggestion to improve test coverage for one of the modified regular expressions to ensure all new formats are validated. Otherwise, the changes are excellent.

[gemini-code-assist]

This change to _VERSION_REGEX correctly adds support for hyphens in the suffix of a full version string (e.g., 1.2.3-ml-debian9). However, this specific format is not covered in the updated test suite. To ensure this change is properly validated, please consider adding a test case for a version string that includes a hyphenated prefix like ml-.