<s> from sqlalchemy import create_engine <EOL> from sqlalchemy . orm import sessionmaker <EOL> from models import Base , CrawlPage <EOL> from data import CrawlSiteData , CrawlSiteListData <EOL> DEFAULT_ENGINE = '<STR_LIT>' <EOL> class CrawlGraphManager ( object ) : <EOL> def __init__ ( self , engine = DEFAULT_ENGINE , autocommit = False , autoflush = False , <EOL> echo = False , drop_all_tables = False , clear_content = False ) : <EOL> self . engine = create_engine ( engine , echo = echo ) <EOL> if drop_all_tables : <EOL> Base . metadata . drop_all ( self . engine ) <EOL> Base . metadata . create_all ( self . engine ) <EOL> self . Session = sessionmaker ( ) <EOL> self . Session . configure ( bind = self . engine , autocommit = autocommit , autoflush = autoflush ) <EOL> self . session = self . Session ( ) <EOL> if clear_content : <EOL> for name , table in Base . metadata . tables . items ( ) : <EOL> self . session . execute ( table . delete ( ) ) <EOL> @ property <EOL> def pages ( self ) : <EOL> return [ page for page in CrawlPage . query ( self . session ) . all ( ) ] <EOL> @ property <EOL> def seeds ( self ) : <EOL> return self . session . query ( CrawlPage ) . filter_by ( is_seed = True ) . all ( ) <EOL> def add_page ( self , url , status = <NUM_LIT:200> , n_redirects = <NUM_LIT:0> , is_seed = False , commit = True ) : <EOL> page , created = CrawlPage . get_or_create ( self . session , url = url ) <EOL> if created : <EOL> page . is_seed = is_seed <EOL> page . status = status <EOL> page . n_redirects = n_redirects <EOL> if commit : <EOL> self . session . commit ( ) <EOL> return page <EOL> def add_link ( self , page , url , commit = True , status = <NUM_LIT:200> ) : <EOL> link_page , created = CrawlPage . get_or_create ( self . session , url = url ) <EOL> if created : <EOL> link_page . status = status <EOL> if link_page not in page . links : <EOL> page . links . append ( link_page ) <EOL> if commit : <EOL> self . session . commit ( ) <EOL> return link_page <EOL> def get_page ( self , url ) : <EOL> return self . session . query ( CrawlPage ) . filter_by ( url = url ) . first ( ) <EOL> def add_site ( self , site , default_status = <NUM_LIT:200> , default_n_redirects = <NUM_LIT:0> ) : <EOL> pages = site . pages if isinstance ( site , CrawlSiteData ) else site <EOL> for i , ( info , links ) in enumerate ( pages ) : <EOL> if isinstance ( info , tuple ) : <EOL> if len ( info ) == <NUM_LIT:2> : <EOL> status , page_url , n_redirects = ( info [ <NUM_LIT:0> ] , info [ <NUM_LIT:1> ] , default_n_redirects ) <EOL> else : <EOL> status , page_url , n_redirects = info <EOL> else : <EOL> status , page_url , n_redirects = ( default_status , info , default_n_redirects ) <EOL> page = self . add_page ( url = page_url , status = status , n_redirects = n_redirects , is_seed = ( i == <NUM_LIT:0> ) ) <EOL> for link_url in links : <EOL> self . add_link ( page = page , url = link_url , status = default_status ) <EOL> def add_site_list ( self , graph , default_status = <NUM_LIT:200> , default_n_redirects = <NUM_LIT:0> ) : <EOL> sites = graph . sites if isinstance ( graph , CrawlSiteListData ) else graph <EOL> for site in sites : <EOL> self . add_site ( site = site , default_status = default_status , default_n_redirects = default_n_redirects ) <EOL> def save ( self ) : <EOL> self . session . commit ( ) <EOL> def render ( self , filename , label = '<STR_LIT>' , labelloc = '<STR_LIT:t>' , labeljust = '<STR_LIT:c>' , <EOL> rankdir = "<STR_LIT>" , ranksep = <NUM_LIT> , <EOL> fontname = '<STR_LIT>' , fontsize = <NUM_LIT> , <EOL> use_urls = False , <EOL> node_fixedsize = '<STR_LIT:true>' , nodesep = <NUM_LIT:0.1> , node_width = <NUM_LIT> , node_height = <NUM_LIT> , node_fontsize = <NUM_LIT:15> , <EOL> include_ids = False ) : <EOL> import pydot <EOL> graph_args = { <EOL> "<STR_LIT>" : rankdir , <EOL> "<STR_LIT>" : ranksep , <EOL> "<STR_LIT>" : nodesep , <EOL> "<STR_LIT>" : fontname , <EOL> "<STR_LIT>" : fontsize , <EOL> } <EOL> if label : <EOL> graph_args . update ( { <EOL> "<STR_LIT>" : labelloc , <EOL> "<STR_LIT>" : labeljust , <EOL> "<STR_LIT:label>" : label <EOL> } ) <EOL> graph = pydot . Dot ( ** graph_args ) <EOL> node_args = { <EOL> "<STR_LIT>" : node_fontsize , <EOL> } <EOL> if use_urls : <EOL> node_seed_shape = '<STR_LIT>' <EOL> node_shape = '<STR_LIT>' <EOL> else : <EOL> node_seed_shape = '<STR_LIT>' <EOL> node_shape = '<STR_LIT>' <EOL> node_args . update ( { <EOL> "<STR_LIT>" : node_fixedsize , <EOL> "<STR_LIT:width>" : node_width , <EOL> "<STR_LIT>" : node_height , <EOL> } ) <EOL> graph . set_node_defaults ( ** node_args ) <EOL> for page in self . pages : <EOL> graph . add_node ( pydot . Node ( name = self . _clean_page_name ( page , include_id = include_ids ) , <EOL> fontname = fontname , <EOL> fontsize = node_fontsize , <EOL> shape = node_seed_shape if page . is_seed else node_shape ) ) <EOL> for link in page . links : <EOL> graph . add_edge ( pydot . Edge ( self . _clean_page_name ( page , include_id = include_ids ) , <EOL> self . _clean_page_name ( link , include_id = include_ids ) ) ) <EOL> graph . write_png ( filename ) <EOL> def _clean_page_name ( self , page , include_id ) : <EOL> cleaned_name = page . url <EOL> cleaned_name = cleaned_name . replace ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> cleaned_name = cleaned_name . replace ( '<STR_LIT>' , '<STR_LIT>' ) <EOL> if include_id : <EOL> cleaned_name = "<STR_LIT>" % ( page . id , cleaned_name ) <EOL> return cleaned_name </s>
