Workstation.py

import csv

def display_menu():
    print("\n=== Computer Equipment Audit Checklist ===")
    print("1. Add a workstation")
    print("2. Add equipment to a workstation")
    print("3. View checklist")
    print("4. Export checklist to CSV")
    print("5. Exit")

def add_workstation(checklist):
    workstation_name = input("Enter workstation name: ").strip()
    if workstation_name in checklist:
        print(f"Workstation '{workstation_name}' already exists.")
    else:
        checklist[workstation_name] = []
        print(f"Workstation '{workstation_name}' added.")

def add_equipment_to_workstation(checklist):
    workstation_name = input("Enter workstation name: ").strip()
    if workstation_name not in checklist:
        print(f"Workstation '{workstation_name}' does not exist. Please add it first.")
        return
    equipment_name = input("Enter equipment name: ").strip()
    if equipment_name in checklist[workstation_name]:
        print(f"'{equipment_name}' is already listed for workstation '{workstation_name}'.")
    else:
        checklist[workstation_name].append(equipment_name)
        print(f"'{equipment_name}' added to workstation '{workstation_name}'.")

def view_checklist(checklist):
    if not checklist:
        print("No workstations have been added yet.")
    else:
        print("\n=== Checklist ===")
        for workstation, equipment in checklist.items():
            print(f"\nWorkstation: {workstation}")
            for item in equipment:
                print(f"  - {item}")

def export_to_csv(checklist):
    if not checklist:
        print("No checklist data to export.")
        return
    filename = input("Enter filename for the CSV (e.g., 'audit_checklist.csv'): ").strip()
    try:
        with open(filename, mode='w', newline='') as file:
            writer = csv.writer(file)
            writer.writerow(["Workstation", "Equipment"])
            for workstation, equipment in checklist.items():
                for item in equipment:
                    writer.writerow([workstation, item])
        print(f"Checklist successfully exported to '{filename}'.")
    except Exception as e:
        print(f"Error exporting to CSV: {e}")

def main():
    checklist = {}
    while True:
        display_menu()
        choice = input("Enter your choice: ").strip()
        if choice == "1":
            add_workstation(checklist)
        elif choice == "2":
            add_equipment_to_workstation(checklist)
        elif choice == "3":
            view_checklist(checklist)
        elif choice == "4":
            export_to_csv(checklist)
        elif choice == "5":
            print("Exiting the program. Goodbye!")
            break
        else:
            print("Invalid choice. Please try again.")

if __name__ == "__main__":
    main()