build-datasets.py

# -*- coding: utf-8 -*-

'''
Copyright (c) Esteban Ponce de León (2023)
Modified by Clarissa R Mendes (claromes)

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
'''

# import modules
import pandas as pd
import argparse
import json
import glob
import time
import os
import tqdm
import telethon

# import local submodules
from telegram_tracker.utils import (
	chats_dataset_columns, clean_msg, msg_attrs, get_forward_attrs, get_reply_attrs,
	get_url_attrs, get_document_attrs, get_poll_attrs, get_contact_attrs,
	get_geo_attrs, msgs_dataset_columns
)

'''

Arguments

'''

parser = argparse.ArgumentParser(description='Arguments.')
parser.add_argument(
	'--data-path',
	'-d',
	type=str,
	required=False,
	help='Path where data is located. Will use `./output/data` if not given.'
)

# Parse arguments
args = vars(parser.parse_args())

# get main path
if args['data_path']:
	main_path = args['data_path']
	if main_path.endswith('/'):
		main_path = main_path[:-1]
else:
	main_path = './output/data'

# log results
text = f'''
Init program at {time.ctime()}

'''
print (text)

# Collect JSON files
json_files_path = f'{main_path}/**/*_messages.json'
json_files = glob.glob(
	os.path.join(json_files_path),
	recursive=True
)

# Collected channels
chats_file_path = f'{main_path}/collected_chats.csv'
data = pd.read_csv(chats_file_path, encoding='utf-8')

# Init values
data['collected_actions'] = 0
data['collected_posts'] = 0
data['replies'] = 0
data['other_actions'] = 0
data['number_views'] = 0
data['forwards'] = 0
data['replies_received'] = 0

# Save dataset
msgs_file_path = f'{main_path}/msgs_dataset.csv'
msgs_data_columns = msgs_dataset_columns()

# JSON files
for f in json_files:
	'''

	Iterate JSON files
	'''
	#  Get channel name
	username = f.split('.json')[0].replace('\\', '/').split('/')[-1].replace(
		'_messages', ''
	)

	# Echo
	print (f'Reading data from channel -> {username}')

	# read JSON file
	with open(f, encoding='utf-8', mode='r') as fl:
		obj = json.load(fl)
		fl.close()

	'''

	Get actions
	'''
	actions = obj['count']
	posts = len(
		[
			i for i in obj['messages'] if 'message' in i.keys()
			and i['reply_to'] == None
		]
	)
	replies = len(
		[
			i for i in obj['messages'] if 'message' in i.keys()
			and i['reply_to'] != None
		]
	)
	other = len(
		[
			i for i in obj['messages'] if 'action' in i.keys()
		]
	)

	'''

	Attrs: views, forwards, replies
	'''
	views = sum(
		[
			i['views'] for i in obj['messages']
			if 'views' in i.keys() and i['views'] != None
		]
	)
	forwards = sum(
		[
			i['forwards'] for i in obj['messages']
			if 'forwards' in i.keys() and i['forwards'] != None
		]
	)
	replies_received = sum(
		[
			i['replies']['replies'] for i in obj['messages']
			if 'replies' in i.keys() and i['replies'] != None
		]
	)

	# Add values to dataset
	data.loc[data['username'] == username, 'collected_actions'] = actions
	data.loc[data['username'] == username, 'collected_posts'] = posts
	data.loc[data['username'] == username, 'replies'] = replies
	data.loc[data['username'] == username, 'other_actions'] = other
	data.loc[data['username'] == username, 'number_views'] = views
	data.loc[data['username'] == username, 'forwards'] = forwards
	data.loc[data['username'] == username, 'replies_received'] = replies_received

	'''

	Reading posts
	'''
	messages = obj['messages']
	pbar = tqdm.tqdm(total=len(messages))
	pbar.set_description(f'Reading posts')

	# main object
	response = {
		'channel_name': username
	}

	for idx, item in enumerate(messages):
		'''

		Iterate posts
		'''
		if item['_'] == 'Message':
			# channel id
			response['channel_id'] = item['peer_id']['channel_id']

			# message id
			msg_id = item['id']
			response['msg_id'] = msg_id

			# add attrs
			response['message'] = item['message']

			# clean message
			msg = clean_msg(item['message'])
			response['cleaned_message'] = msg

			# timestamp
			date = item['date']
			response['date'] = date

			# signature and message link
			response['signature'] = \
				f'msg_iteration.{idx}.user.{username}.post.{msg_id}'
			response['msg_link'] = f'https://t.me/{username}/{msg_id}'

			# check peer
			response['msg_from_peer'] = None
			response['msg_from_id'] = None
			response = msg_attrs(item, response)

			# reactions
			response['views'] = 0 if item['views'] == None else item['views']
			response['number_replies'] = \
				item['replies']['replies'] if item['replies'] != None else 0
			response['number_forwards'] = 0 if item['forwards'] == None \
				else item['forwards']

			# Forward attrs
			forward_attrs = item['fwd_from']
			response['is_forward'] = 1 if forward_attrs != None else 0

			response['forward_msg_from_peer_type'] = None
			response['forward_msg_from_peer_id'] = None
			response['forward_msg_from_peer_name'] = None
			response['forward_msg_date'] = None
			response['forward_msg_date_string'] = None
			response['forward_msg_link'] = None
			if forward_attrs:
				response = get_forward_attrs(
					forward_attrs,
					response,
					data
				)

			# Reply attrs
			response['is_reply'] = 0
			response['reply_to_msg_id'] = None
			response['reply_msg_link'] = None
			response = get_reply_attrs(
				item,
				response,
				username
			)

			# Media
			response['contains_media'] = 1 if item['media'] != None else 0
			if 'media' in item.keys():
				response['media_type'] = None if item['media'] == None \
					else item['media']['_']

			# URLs -> Constructor MessageMediaWebPage
			'''
			Type WebPage

			Source: https://core.telegram.org/constructor/messageMediaWebPage
			Telethon: https://tl.telethon.dev/constructors/web_page.html
			'''
			response = get_url_attrs(item['media'], response)

			# Media Document -> Constructor MessageMediaDocument
			'''
			Type Document

			Source: https://core.telegram.org/constructor/messageMediaDocument
			Telethon: https://tl.telethon.dev/constructors/document.html
			'''
			response['document_type'] = None
			response['document_id'] = None
			response['document_video_duration'] = None
			response['document_filename'] = None

			response = get_document_attrs(item['media'], response)


			# Polls attrs
			'''

			Type Poll

			Source: https://core.telegram.org/constructor/messageMediaPoll
			Telethon: https://tl.telethon.dev/constructors/poll.html

			'''
			response['poll_id'] = None
			response['poll_question'] = None
			response['poll_total_voters'] = None
			response['poll_results'] = None
			response = get_poll_attrs(item['media'], response)

			# Contact attrs
			'''

			Type Contact

			Source: https://core.telegram.org/constructor/messageMediaContact
			Telethon: https://tl.telethon.dev/constructors/message_media_contact.html
			'''
			response['contact_phone_number'] = None
			response['contact_name'] = None
			response['contact_userid'] = None
			response = get_contact_attrs(item['media'], response)

			# Geo attrs
			'''

			Type GeoPoint

			Source: https://core.telegram.org/constructor/messageMediaGeo
			Telethon:
			>	https://tl.telethon.dev/constructors/geo_point.html
			>	https://tl.telethon.dev/constructors/message_media_venue.html

			'''
			response['geo_type'] = None
			response['lat'] = None
			response['lng'] = None
			response['venue_id'] = None
			response['venue_type'] = None
			response['venue_title'] = None
			response['venue_address'] = None
			response['venue_provider'] = None
			response = get_geo_attrs(item['media'], response)

			# create dataframe
			res = [response]
			df = pd.DataFrame.from_dict(res)

			# order df msgs data columns
			df = df[msgs_data_columns].copy()

			# update CSV file
			df.to_csv(
				msgs_file_path,
				encoding='utf-8',
				header=not os.path.isfile(msgs_file_path),
				index=False,
				mode='a'
			)

		# Update pbar
		pbar.update(1)

	# Close pbar connection
	pbar.close()

	print ('-- END --')
	print ('')

# Save data
chats_columns = chats_dataset_columns()
data = data[chats_columns].copy()

data.to_excel(
	chats_file_path.replace('.csv', '.xlsx'),
	index=False
)