Welcome to the Treehouse Community

Want to collaborate on code errors? Have bugs you need feedback on? Looking for an extra set of eyes on your latest project? Get support with fellow developers, designers, and programmers of all backgrounds and skill levels here with the Treehouse Community! While you're at it, check out some resources Treehouse students have shared here.

Looking to learn something new?

Treehouse offers a seven day free trial for new students. Get access to thousands of hours of content and join thousands of Treehouse students and alumni in the community today.

Start your free trial

Python

Posted by Aaron West
Aaron West
256 Points

Select the input value of a specific argument (argeparse) to use a class object.

working with the following code:

def get_args():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description=textwrap.dedent('''\
            The Easy PSQL CLI tool
            --------------------------------
            Easily edit your postgres configuration,
            Add users, and create databases.
            * Uses default postgres user *
            '''))
    parser._positionals.title = 'Positional arguments'
    parser._optionals.title = 'Optional arguments'
    parser.add_argument("--add-user",
                        required=False,
                        action='store',
                        default='',
                        help="Add a new PostgrSQL user")
    parser.add_argument("--del-user",
                        required=False,
                        action='store',
                        default='',
                        help="Delete a PostgrSQL user")
    parser.add_argument("--passwd",
                        required=False,
                        action='store',
                        default='',
                        help="Assign the new PostgreSQL user a password")
    args = parser.parse_args()
    args_dict = vars(args)

I need to take the input value of any of these arguments so that I can use them within a few objects later on. Preferably i'd like to turn them into variables, for example, make --add-user = dbuser as a variable.

On top of this, I would like to be able to say if the argument --add-user was used do this.. else --del-user do this, for example (some none-working code):

    def dbuser(self):
        if get_args.args('add-user'):
            self.cur.execute(f"create user {self.user} with password '{self.password}'")
        if get_args.args('del-user'):
            self.cur.execute(f"drop user {self.user} ")

2 Answers

Chris Freeman
MOD
Chris Freeman
Treehouse Moderator 68,460 Points
Chris Freeman
Chris Freeman
Treehouse Moderator 68,460 Points

Your usage model needs a few tweaks:

import argparse
import textwrap


def get_args(arglist):  # Add 'arglist' parameter to receive arguments passed in

    # parser object set up looks ok....
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description=textwrap.dedent('''\
            The Easy PSQL CLI tool
            --------------------------------
            Easily edit your postgres configuration,
            Add users, and create databases.
            * Uses default postgres user *
            '''))
    parser._positionals.title = 'Positional arguments'
    parser._optionals.title = 'Optional arguments'
    parser.add_argument("--add-user",
                        required=False,
                        action='store',
                        default='',
                        help="Add a new PostgrSQL user")
    parser.add_argument("--del-user",
                        required=False,
                        action='store',
                        default='',
                        help="Delete a PostgrSQL user")
    parser.add_argument("--passwd",
                        required=False,
                        action='store',
                        default='',
                        help="Assign the new PostgreSQL user a password")

    # need to pass argument list into the parse
    args = parser.parse_args(arglist)
    args_dict = vars(args)  # this may not be needed. See below
    return args, args_dict  # need to return the parsed arguments


# calling get_args()
args, arg_dict = get_args(["--add-user", "foobar", "--del-user", "snafu"])

# this produces a Namspace that allows accessing the arguments and their values
print(args)
# arg_dict for reference
print(arg_dict)

if args.add_user:
    print(f"do something to add user {args.add_user}")
if args.del_user:
    print(f"do something to delete user {args.del_user}")

Saved to a file and run produces:

$ python arg_parser.py --add-user foobar --del-user foobat
Namespace(add_user='foobar', del_user='snafu', passwd='')
{'add_user': 'foobar', 'del_user': 'snafu', 'passwd': ''}
do something to add user foobar
do something to delete user snafu

As you can see, the returned args object has what you need, removing the use of args_dict.

Also, consider:

  • moving the parser definition outside of the get_args() function. So it isn't rebuilt every time you call get_args()
  • alternatively, calling get_args() once, then re-referencing the args object as needed.

Post back if you need more help. Good Luck!!

Aaron West
Aaron West
256 Points

so with args, arg_dict = get_args(['--add-user', 'foobar']), how would i accept user CLI input such as the following: ./myscript --add-user myuser as the input, in placement of foobar?

Also, I don't understand the need for arglist in get_args(arglist) since it only seems to make it so that I must define the input of get_args within the script now, as opposed to an external CLI, which is what I am trying to do. arglist seems to actually enforce the use of args, arg_dict = get_args([]) whereas before using arglist, it did not.

Chris Freeman
Chris Freeman
Treehouse Moderator 68,460 Points
Chris Freeman
Chris Freeman
Treehouse Moderator 68,460 Points

Sorry, the "save and run" portion of my post was for a different solution.

Below is an example of how I would use it from the command line. The if __name__ == "__main__" and the KeyboardInterrupt help interactive use by gracefully capturing a Ctrl-C to kill the program.

#!/usr/bin/env python
# The above line allows my virtual env to set the version of python used  at runtime
import argparse
import textwrap
import sys

# helper functions
def parse_args(args):
    # parser object set up looks ok....
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description=textwrap.dedent('''\
            The Easy PSQL CLI tool
            --------------------------------
            Easily edit your postgres configuration,
            Add users, and create databases.
            * Uses default postgres user *
            '''))
    parser._positionals.title = 'Positional arguments'
    parser._optionals.title = 'Optional arguments'
    parser.add_argument("--add-user",
                        required=False,
                        action='store',
                        default='',
                        help="Add a new PostgrSQL user")
    parser.add_argument("--del-user",
                        required=False,
                        action='store',
                        default='',
                        help="Delete a PostgrSQL user")
    parser.add_argument("--passwd",
                        required=False,
                        action='store',
                        default='',
                        help="Assign the new PostgreSQL user a password")
    return parser.parse_args(args)


def guts_of_program(args):
    if args.add_user:
        print(f"do something to add user {args.add_user}")

    if args.del_user:
        print(f"do something to delete user {args.del_user}")

    status_or_results = True  # or something real

    return status_or_results


# main program
def main():
    args = parse_args(sys.argv[1:])
    results = guts_of_program(args)
    print(results)


if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        print("Keyboard Interrupt requested...exiting")
        sys.exit(0)

I saved the file as arg_parser_sys.py then ran the file in the session below:

(tth_forum) chris@chris-pop-os:~/devel/Treehouse/forum/aaron_west
$ python arg_parser_sys.py 
True
(tth_forum) chris@chris-pop-os:~/devel/Treehouse/forum/aaron_west
$ python arg_parser_sys.py --add-user good_user --del-user bad_user
do something to add user good_user
do something to delete user bad_user
True

# Changing Linux mode to allow direct execution.
(tth_forum) chris@chris-pop-os:~/devel/Treehouse/forum/aaron_west
$ chmod +x arg_parser_sys.py
(tth_forum) chris@chris-pop-os:~/devel/Treehouse/forum/aaron_west
$ ./arg_parser_sys.py --add-user good_user --del-user bad_user
do something to add user good_user
do something to delete user bad_user
True
(tth_forum) chris@chris-pop-os:~/devel/Treehouse/forum/aaron_west
$

Hope this helps. Good luck!!

Aaron West
Aaron West
256 Points
Aaron West
Aaron West
256 Points

Chris Freeman - My code that I came up with is as follows:

#!/usr/bin/env python

import sys
import argparse
import textwrap
import psycopg2
from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT

####################
'''INITIALIZE'''
####################

def parse_args(args):
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description=textwrap.dedent('''
            The Easy PSQL CLI tool
            --------------------------------
            Easily edit your postgres configuration,
            Add users, and create databases.
            * Uses default postgres user *
            '''))
    parser._positionals.title = 'Positional arguments'
    parser._optionals.title = 'Optional arguments'
    parser.add_argument('object',
                    default='',
                    help='USERNAME, DATABASE, or OBJECT name to be altered, created, or deleted (default: %(default)s)')
    parser.add_argument('--add',
                    default='',
                    nargs='?',
                    choices=['user', 'database'],
                    help='Add user or database (default: %(default)s)')
    parser.add_argument('--delete',
                    default='',
                    nargs='?',
                    choices=['user', 'database'],
                    help='Add user or database (default: %(default)s)')
    parser.add_argument("-W", "--password",
                    required=False,
                    action='store',
                    default='',
                    help="Assign the new PostgreSQL user a password")
    return parser.parse_args(args)

class PgConnect:

    def __init__(self, host='127.0.0.1', dbname='', user='postgres', password=''):
        self.host = host
        self.dbname = dbname
        self.user = user
        self.password = password
        self.con = None

    def connectdb(self):
        self.con = psycopg2.connect(
            f"host='{self.host}' "
            f"user='{self.user}' "
            f"password='{self.password}'")
        self.con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
        self.cur = self.con.cursor()

    def usermod(self):
        if args.add == 'user':
            self.cur.execute(f"create user {args.object} with password '{args.password}' ")
        if args.delete == 'user':
            self.cur.execute(f"drop user {args.object} ")

    def dbmod(self):
        if args.add == 'database':
            self.cur.execute(f"create database {args.object} ")
        if args.delete == 'database':
            self.cur.execute(f"drop database {args.object} ")

    def disconnect(self):
        self.cur.close()

args = parse_args(sys.argv[1:])

###############
'''BODY'''
##############

con = PgConnect()
con.connectdb()
con.usermod()
con.dbmod()

###############
'''TESTING'''
###############

'''
if args.add == 'user':
    print(f"do something to add user {args.object}")
if args.add == 'database':
    print(f"do something to add database {args.object}")
'''

Thank you for your help! Is there anything that you would suggest I add (such as additional logic) to make it a more professional body of code, perhaps best practice recomendations if I am missing anything?

Chris Freeman
Chris Freeman
Treehouse Moderator 68,460 Points
Chris Freeman
Chris Freeman
Treehouse Moderator 68,460 Points

I would pass args or specific args values like args.add as arguments to the PgConnect methods instead of relying on args as a global variable. The rest looks OK. Beyond that, it depends on how any additional code is implemented.

Consider what would happen if your code was imported by another module. The args = parse_args(sys.argv[1:]) would fail since there would be any sys.argv available. That is why I put the working commands in a main() routine and use the if __name__ == "__main__" to only run the main() code if the file is run interactively from the command line and will not run main() if the file is imported into another file.