code-generator/code_generator.py at main · MariosGramm/code-generator · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
from argparse import ArgumentParser
from code_gen import generate_code as code_generation
from code_gen import save_code


def generate_code(description: str, with_tests: bool = False, filename: str = None):
    code, tests = code_generation(description)

    if filename:
        save_code(code, filename)

    if with_tests:
        return code, tests
    else:
        return code, None

def interactive_mode():

    print("""
============================================================================================
||                                                                                        ||
||   ██████╗ ██████╗ ██████╗ ███████╗                                                     ||
||  ██╔════╝██╔═══██╗██╔══██╗██╔════╝                                                     ||
||  ██║     ██║   ██║██║  ██║█████╗                                                       ||
||  ██║     ██║   ██║██║  ██║██╔══╝                                                       ||
||  ╚██████╗╚██████╔╝██████╔╝███████╗                                                     ||
||   ╚═════╝ ╚═════╝ ╚═════╝ ╚══════╝                                                     ||
||                                                                                        ||
||   ██████╗ ███████╗███╗   ██╗███████╗██████╗  █████╗ ████████╗ ██████╗ ██████╗          ||
||  ██╔════╝ ██╔════╝████╗  ██║██╔════╝██╔══██╗██╔══██╗╚══██╔══╝██╔═══██╗██╔══██╗         ||
||  ██║  ███╗█████╗  ██╔██╗ ██║█████╗  ██████╔╝███████║   ██║   ██║   ██║██████╔╝         ||
||  ██║   ██║██╔══╝  ██║╚██╗██║██╔══╝  ██╔══██╗██╔══██║   ██║   ██║   ██║██╔══██╗         ||
||  ╚██████╔╝███████╗██║ ╚████║███████╗██║  ██║██║  ██║   ██║   ╚██████╔╝██║  ██║         ||
||   ╚═════╝ ╚══════╝╚═╝  ╚═══╝╚══════╝╚═╝  ╚═╝╚═╝  ╚═╝   ╚═╝    ╚═════╝ ╚═╝  ╚═╝         ||
||                                                                                        ||
============================================================================================
""")

    print("")
    print("Welcome to the code generator!")


    while True:
        print("")
        description = input("Please provide a code description: ")

        if description:
            description = description.strip()
        else:
            print("Code decription is empty. Please try again.")
            print("")
            continue


        with_tests_input = ""

        while True:
            with_tests_input = input("Would you like tests to be generated?(y/n): ")

            if with_tests_input:
                with_tests_input = with_tests_input.strip()
            if with_tests_input.lower() == "y":
                with_tests = True
                break
            elif with_tests_input.lower() == "n":
                with_tests = False
                break
            print("")
            print("Please provide a valid answer (y/n) or (Y/N)")

        while True:
            print("")
            print("Please provide a filename for the code to be saved at.")
            print("If you don't want the code to be saved, just continue by pressing enter.")
            print("")

            filename = input("Filename: ")

            if filename:
                filename = filename.strip()
            else:
                break

            forbidden_character_found = False
            for char in ['\\', '/', ':', '*', '?', '"', '<', '>', '|']:
                if char in filename:
                    print(f"Filename containts a forbidden character: '{char}'")
                    print("Please try again. File cannot be saved.")
                    forbidden_character_found = True
                    break

            if forbidden_character_found:
                continue
            else:
                break

        code, tests = generate_code(description, with_tests, filename)

        print("\n-- Generated Code --\n")
        print(code)

        if tests:
            print("\n-- Generated Tests --\n")
            print(tests)

        while True:
            print("")
            again = input("Would you like to continue generating code? (y/n): ")

            if again:
                again = again.strip().lower()
            if again.startswith("y"):
                break
            elif again.startswith("n"):
                print("Exit...")
                return
            else:
                print("")
                print("Please provide a valid answer (y/n) or (Y/N)")


def main():
    parser = ArgumentParser(prog= "Python Code Generator",
                            description= "Python Code Generator from natural language")
    parser.add_argument("description", nargs="?", help="Short description of the code that will be generated")
    parser.add_argument("--with-tests", action="store_true", help= "Generate tests for the generated code")
    parser.add_argument("--save", metavar= "FILE", help= "Filename to save the generated code")
    parser.add_argument("-i", "--interactive", action= "store_true", help= "Run in interactive mode")

    args = parser.parse_args()

    if args.interactive or not args.description:
        interactive_mode()
        return

    code, tests = generate_code(
        description= args.description,
        with_tests = args.with_tests,
        filename= args.save
    )

    print("\n-- Generated Code --\n")
    print(code)

    if args.with_tests and tests:
        print("\n-- Generated Tests --\n")
        print(tests)


if __name__ == "__main__":
    main()