Usuário:Danilo.mac/buscadump.py

Origem: Wikipédia, a enciclopédia livre.

Script para busca no dump de artigos do Tool Labs. Modificando as variáveis 'ddir' e 'dump' pode funcionar também em um dump local. Este script não utiliza o pywikipediabot.


"""
@ Autor:     [[Usuário:Danilo.mac]]

@ Licença:   GNU General Public License 3.0 (GPL V3) e Creative Commons Attribution/Share-Alike (CC-BY-SA)

Descrição:   Script para contagem e busca no último dump dos artigos da
             Wikipédia lusófona.

Uso:         python buscadump.py [-c] termo_ou_regex_a_ser_buscado [arquivo_onde_salvar]

             Este script funciona exclusivamente no TooL Labs, não é necessário baixar os
             dumps, o script faz a busca na última cópia do dump de artigos que é mantido
             no Tool Labs. Ao utilizar no modo busca, a lista de artigos é salva em um
             arquivo passado como argumento ou em 'busca.txt' se o arquivo não for
             especificado.
             
             A busca é feita analisando 50.000.000 bytes por vez, que corresponde a
             aproximadamente 1% do dump. Cada vez que essa quantiade é analisada, o script
             pergunta se deseja continuar, a pergunta é respondida sim automaticamente em
             alguns segundos.

Opções:
-c           Faz a busca no modo contagem, apenas exibindo o número de artigos que contém
             correspondências para o termo buscado. Nenhum arquivo é salvo neste modo.
"""
import bz2, re, signal, codecs, sys, os

regexPag = re.compile(ur'(?s)<title>([^\n]+?)</title>\s*<ns>0</ns>.*?<text xml:space="preserve">(.*?)</text>')
regexRedir = re.compile(ur'(?i)#redirec')
subst = lambda texto: re.sub(r'\s+', '', texto.replace('&lt;', '<').replace('&gt;', '>').replace('&quot;', '"').replace('&amp;', '&'))
ddir = '/public/datasets/public/ptwiki/' + max(os.listdir('/public/datasets/public/ptwiki/'))
dump = ddir + '/' + [a for a in os.listdir(ddir) if a.endswith('-pages-articles.xml.bz2')][0]

def buscar(busca, salvar=None, regex=False):
  f = bz2.BZ2File(dump)
  busca = regex and re.compile(busca.decode('utf-8')) or busca
  continuar = True
  titulo = None
  lista = []
  c = 0
  temp = f.read(50000000)
  while True:
    for pag in regexPag.finditer(temp):
      if len(pag.group(2)) < 1000 and regexRedir.search(pag.group(2)):
        continue
      txt = pag.group(2).replace('&gt;', '>').replace('&lt;', '<')
      if not regex and busca in subst(txt):
        lista.append(pag.group(1).decode('utf-8'))
      elif regex:
        termos = busca.findall(subst(txt).decode('utf-8'))
        if termos:
          lista.append(u'{} -> {}'.format(pag.group(1).decode('utf-8'), u'; '.join(set(termos))))
      c += 1
    if '</mediawiki>' in temp:
      break
    elif timer_input(str(len(lista)) + ' em ' + str(c) + ', continuar?(s/n) ', 4, 's') == 's':
      if sum(map(len, lista)) > 1000000:
       r = raw_input('O resultado tem mais de 1MB, continuar mesmo assim?([c]ontinuar/parar e [s]alvar/parar e [n]ão salvar)')
       if r == 's':
         break
       elif r == 'n':
         return
      temp = temp[temp.rfind('<page>'):] + f.read(50000000)
    else:
      print 'verificados', c, 'artigos'
      break
  f.close()
  if salvar and lista:
    with codecs.open(salvar,'w','utf8') as s:
      s.write(u'\n'.join(lista))
    print u'Salvo {} títulos em {}'.format(len(lista), salvar)
  else:
    return lista

def contar(busca, regex=False):
  print dump
  f = bz2.BZ2File(dump)
  busca = regex and re.compile(busca.decode('utf-8')) or busca
  continuar = True
  cont = 0
  c = 0
  temp = f.read(50000000)
  while True:
    for pag in regexPag.finditer(temp):
      if len(pag.group(2)) < 1000 and regexRedir.search(pag.group(2)):
        continue
      txt = pag.group(2).replace('&gt;', '>').replace('&lt;', '<')
      if not regex and busca in subst(txt):
        cont += 1
      elif regex and busca.search(txt.decode('utf-8')):
        cont += 1
      c += 1
    if '</mediawiki>' in temp:
      break
    elif timer_input(str(cont) + ' em ' + str(c) + ', continuar?(s/n) ', 3, 's') == 's':
      temp = temp[temp.rfind('<page>'):] + f.read(50000000)
    else:
      break
  f.close()
  print cont, 'em', c

def timer_input(text='?', time=10, resp=''):
  class Parar: pass
  def parar(signum, frame): raise Parar
  signal.signal(signal.SIGALRM, parar)
  signal.alarm(time)
  try:
   r = raw_input(text)
   signal.alarm(0)
   return r
  except Parar: print 'respondendo \'%s\'' % resp
  signal.signal(signal.SIGALRM, signal.SIG_IGN)
  return resp

if __name__ == "__main__":
  args = sys.argv[1:]
  if len(args) == 2 and args[0] == '-c':
    if raw_input('Contagem com regex?(s/n) ') == 's':
      contar(args[1], True)
    else:
      contar(args[1])
  elif len(args) in (1, 2):
    if len(args) == 1:
      args.append('busca.txt')
    if raw_input('Busca com regex?(s/n) ') == 's':
      args.append(True)
    buscar(*args)
  else:
    print u'Uso: python buscadump.py [-c] termo_a_ser_buscado [arquivo_onde_salvar]'