Сгруппировать и найти среднее по словарю

import json
def age_count(x):
    s = ('{ "workers": [{"name": "Имя1", "age": 34, "occupation": "Программист"}, {"name": "Имя2", "age": 24, "occupation": "Программист"}, {"name": "Имя3", "age": 40, "occupation": "Грузчик"} ] }')
    
    
    r = {}
    for x in json.loads(s)['workers'] :
    
    r.setdefault(x["occupation"],[]).append(x["age"])
    for x in r.keys() :
    
    r[x] = sum(r[x])/len(r[x])
    print(r)
    
    
with open("input.txt") as f:
    age_count(f.read())

Условия:
Напишите функцию age_count, которая принимает JSON строку следующего вида:

{ "workers": [{"name": "Имя1", "age": 34, "occupation": "Программист"}, {"name": "Имя2", "age": 24, "occupation": "Программист"}, {"name": "Имя3", "age": 40, "occupation": "Грузчик"} ] }

Далее, в функции нужно посчитать средний возраст по каждой из профессий и сохранить в файл age.json этот возраст в следующем виде:

{"Программист": 29.0, "Грузчик": 40.0}

Также результирующий JSON должен быть выведен на экран в конце выполнения функции.

Для проверки задания после описания функции добавьте следующие строки:

with open("input.txt") as f:
    age_count(f.read())

Формат ввода (содержимое input.txt)

{
    "workers": [
         {
            "name": "Имя1",
            "age": 34,
            "occupation": "Программист"
        }, {
            "name": "Имя2",
            "age": 24,
            "occupation": "Программист"
        }, {
            "name": "Имя3",
            "age": 40,
            "occupation": "Грузчик"
        }
    ]
}

Формат вывода:

{"Программист": 29.0, "Грузчик": 40.0}

from itertools import groupby

def avg(lst):
    return sum(lst) / len(lst)

data = json.loads(s)

res = {a:avg([z["age"] for z in b]) 
       for a,b in groupby(data['workers'], key=lambda x: x['occupation'])}

In [67]: res
Out[67]: {'Программист': 29.0, 'Грузчик': 40.0}

import json


def age_count(data):
    s = json.loads(data)

    r = dict()
    employees = s['workers']

    for employee in employees:
        if employee['occupation'] not in r:
            r[employee['occupation']] = {
                'summary_age': 0,
                'employees_qty': 0
            }
        r[employee['occupation']]['summary_age'] += employee['age']
        r[employee['occupation']]['employees_qty'] += 1

    json.dump(
        {
            occupation: r[occupation]['summary_age'] / r[occupation]['employees_qty'] for occupation in r
        }, 
        open('age.json', 'w', encoding='utf-8'),
        ensure_ascii=False
    )


with open("input.txt", 'r', encoding='utf-8') as f:
    age_count(f.read())

check = json.load(open('age.json', 'r', encoding='utf-8'))
print(check)
# {'Программист': 29.0, 'Грузчик': 40.0}

{
    "workers": [{
            "name": "Имя1",
            "age": 34,
            "occupation": "Программист"
        }, {
            "name": "Имя2",
            "age": 24,
            "occupation": "Программист"
        }, {
            "name": "Имя3",
            "age": 40,
            "occupation": "Грузчик"
        }
    ]
}